Mojo function
foreach
foreach[type: DType, rank: Int, //, func: fn[Int](Index[rank]) capturing -> SIMD[type, $0], *, target: StringSlice[StaticConstantOrigin] = __init__[__mlir_type.!kgen.string]("cpu"), simd_width: Int = get_kernel_simd_width[::DType,::StringSlice[::Bool(), _synchronous: Bool = False, _trace_name: StringSlice[StaticConstantOrigin] = __init__[__mlir_type.!kgen.string]("mogg.for_each")](tensor: ManagedTensorSlice[io_spec, static_spec=static_spec], ctx: DeviceContextPtr = DeviceContextPtr())
Apply the function func
to each element of the tensor slice.
Parameters:
- type (
DType
): The data type of the elements in the tensor slice. - rank (
Int
): The rank of the tensor slice. - func (
fn[Int](Index[rank]) capturing -> SIMD[type, $0]
): The function to apply to each element of the tensor slice. - target (
StringSlice[StaticConstantOrigin]
): Indicates the type of the target device (e.g. "cpu", "gpu"). - simd_width (
Int
): The SIMD width for the target (usually leave this as its default value). - _synchronous (
Bool
): True to run the custom op synchronously in the runtime (defaults to False). - _trace_name (
StringSlice[StaticConstantOrigin]
): Name of the executed operation displayed in the trace_description.
Args:
- tensor (
ManagedTensorSlice[io_spec, static_spec=static_spec]
): The output tensor slice which receives the return values fromfunc
. - ctx (
DeviceContextPtr
): The call context (forward this from the custom operation).
foreach[: origin.set, type: DType, rank: Int, //, func: fn[Int](Index[rank]) capturing -> SIMD[type, $0], out_func: fn[Int, Int](Index[$1]) capturing -> None, *, target: StringSlice[StaticConstantOrigin] = __init__[__mlir_type.!kgen.string]("cpu"), simd_width: Int = get_kernel_simd_width[::DType,::StringSlice[::Bool(), _synchronous: Bool = False, _trace_name: StringSlice[StaticConstantOrigin] = __init__[__mlir_type.!kgen.string]("mogg.for_each")](tensor: ManagedTensorSlice[io_spec, static_spec=static_spec], ctx: DeviceContextPtr = DeviceContextPtr())
Apply the function func
to each element of the tensor slice.
Parameters:
- type (
DType
): The data type of the elements in the tensor slice. - rank (
Int
): The rank of the tensor slice. - func (
fn[Int](Index[rank]) capturing -> SIMD[type, $0]
): The function to apply to each element of the tensor slice. - out_func (
fn[Int, Int](Index[$1]) capturing -> None
): The function to apply on each output element. - target (
StringSlice[StaticConstantOrigin]
): Indicates the type of the target device (e.g. "cpu", "gpu"). - simd_width (
Int
): The SIMD width for the target (usually leave this as its default value). - _synchronous (
Bool
): True to run the custom op synchronously in the runtime (defaults to False). - _trace_name (
StringSlice[StaticConstantOrigin]
): Name of the executed operation displayed in the trace_description.
Args:
- tensor (
ManagedTensorSlice[io_spec, static_spec=static_spec]
): The input tensor slice which the consumed values. - ctx (
DeviceContextPtr
): The call context (forward this from the custom operation).
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!