``for_each``: apply a function in-place over a tiled tensor on an AIE core.
| iron.algorithms.for_each.for_each |
( |
|
func, |
|
|
|
tensor, |
|
|
* |
params, |
|
|
|
tile_size = 16 |
|
) |
| |
In-place transform. Internally uses separate input/output ObjectFifos,
but fills and drains to same tensor.
Args:
func: Function to apply, either a lambda/callable or ExternalFunction.
For ExternalFunction, arg_types should be [input_tile, output_tile, *params]
tensor: The tensor to apply in-place transformation
*params: Additional parameters for ExternalFunction only.
Scalar dtypes (np.int32, etc.) are passed as MLIR constants;
array types are transferred via ObjectFifos.
tile_size: Size of each tile processed by a worker (default: 16)
Example::
# kernel has separate in/out tile buffers, but only one tensor is passed
scale = ExternalFunction("scale", arg_types=[tile_ty, tile_ty, scalar_ty, np.int32], ...)
for_each(scale, tensor, factor, tile_size=16)
Returns:
mlir.ir.Module: The compiled MLIR module ready for execution.