IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /docs/manual/basics.md). For the complete Mojo documentation index, see llms.txt.
Skip to main content
Version: 1.0.0b2
For the complete Mojo documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /docs/manual/basics.md).

elementwise

def elementwise[func: def[width: Int, alignment: Int = 1](Coord[_]) capturing -> None, simd_width: Int, *, target: StringSlice[StaticConstantOrigin] = StringSlice("cpu"), _trace_description: StringSlice[StaticConstantOrigin] = StringSlice("elementwise")](shape: Int, context: DeviceContext)

Executes func[width, rank](indices), possibly as sub-tasks, for a suitable combination of width and indices so as to cover shape. Returns when all sub-tasks have completed.

Parameters:

  • ​func (def[width: Int, alignment: Int = 1](Coord[_]) capturing -> None): The body function.
  • ​simd_width (Int): The SIMD vector width to use.
  • ​target (StringSlice[StaticConstantOrigin]): The target to run on.
  • ​_trace_description (StringSlice[StaticConstantOrigin]): Description of the trace.

Args:

  • ​shape (Int): The shape of the buffer.
  • ​context (DeviceContext): The device context to use.

Raises:

If the operation fails.

def elementwise[func: def[width: Int, alignment: Int = 1](Coord[_]) capturing -> None, simd_width: Int, *, target: StringSlice[StaticConstantOrigin] = StringSlice("cpu"), _trace_description: StringSlice[StaticConstantOrigin] = StringSlice("elementwise")](shape: Coord, context: DeviceContext)

Executes func[width, rank](indices), possibly as sub-tasks, for a suitable combination of width and indices so as to cover shape. Returns when all sub-tasks have completed.

Parameters:

  • ​func (def[width: Int, alignment: Int = 1](Coord[_]) capturing -> None): The body function.
  • ​simd_width (Int): The SIMD vector width to use.
  • ​target (StringSlice[StaticConstantOrigin]): The target to run on.
  • ​_trace_description (StringSlice[StaticConstantOrigin]): Description of the trace.

Args:

  • ​shape (Coord): The shape of the buffer.
  • ​context (DeviceContext): The device context to use.

Raises:

If the operation fails.

def elementwise[FuncType: ImplicitlyCopyable & RegisterPassable & def[width: Int, alignment: Int, .element_types.values2x: KGENParamList[CoordLike], .element_types`2x1: TypeList[values]](Coord[element_types]) -> None, simd_width: Int, *, target: StringSlice[StaticConstantOrigin] = StringSlice("cpu"), _trace_description: StringSlice[StaticConstantOrigin] = StringSlice("elementwise")](func: FuncType, shape: Coord, context: DeviceContext)``

Unified-closure entry point for elementwise (DeviceContext).

Accepts a parametric body (already in unified-closure form, with explicit captures) and dispatches to _elementwise_impl. rank and FuncType are inferred from the runtime shape and func arguments, so simd_width is the only explicit positional parameter — callers can write elementwise[N](func, shape, ctx).

Parameters:

  • ​FuncType (ImplicitlyCopyable & RegisterPassable & def[width: Int, alignment: Int, .element_types.values``2x: KGENParamList[CoordLike], .element_types2x1: TypeList[values]](Coord[element_types]) -> None): A parametric callable taking IndexList[rank]and template parameterswidth, rank, alignment`.
  • ​simd_width (Int): The SIMD vector width to use.
  • ​target (StringSlice[StaticConstantOrigin]): The target to run on.
  • ​_trace_description (StringSlice[StaticConstantOrigin]): Description of the trace.

Args:

  • ​func (FuncType): The body closure value.
  • ​shape (Coord): The shape of the buffer.
  • ​context (DeviceContext): The device context to use.

Raises:

If the operation fails.

View source