Skip to main content
Log in

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 from func.
  • 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).