For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo module

functional

Implements higher-order functions.

You can import these APIs from the algorithm package. For example:

from max.algorithm import elementwise

comptime values

stencil

comptime stencil[shape_element_type: DType, input_shape_element_type: DType, //, rank: Int, stencil_rank: Int, stencil_axis: IndexList[stencil_rank, element_type=element_type], simd_width: Int, dtype: DType, map_fn: def(IndexList[stencil_rank, element_type=*?]) -> Tuple[IndexList[stencil_rank], IndexList[stencil_rank]], map_strides: def(dim: Int) -> Int, load_fn: def[simd_width: Int, dtype: DType](IndexList[rank, element_type=*?]) -> SIMD[dtype, simd_width], compute_init_fn: def[simd_width: Int]() -> SIMD[dtype, simd_width], compute_fn: def[simd_width: SIMDLength](IndexList[rank, element_type=*?], SIMD[dtype, simd_width], SIMD[dtype, simd_width]) -> SIMD[dtype, simd_width], compute_finalize_fn: def[simd_width: SIMDLength](IndexList[rank, element_type=*?], SIMD[dtype, simd_width]) -> None] where (eq map_fn.stencil_rank, stencil_rank) where (eq load_fn.rank, rank) where (eq compute_init_fn.dtype, dtype) where (eq compute_fn.rank, rank) where (eq compute_fn.dtype, dtype) where (eq compute_finalize_fn.rank, rank) where (eq compute_finalize_fn.dtype, dtype) = fn_literal

Computes stencil operation in parallel.

Computes output as a function that processes input stencils, stencils are computed as a continuous region for each output point that is determined by map_fn : map_fn(y) -> lower_bound, upper_bound. The boundary conditions for regions that fail out of the input domain are handled by load_fn.

Args: shape: The shape of the output buffer. input_shape: The shape of the input buffer. map_fn_closure: Closure mapping output points to input co-domain bounds. map_strides_closure: Closure returning the stride for a given dimension. load_fn_closure: Closure loading a SIMD vector from input. compute_init_fn_closure: Closure initializing the stencil accumulator. compute_fn_closure: Closure processing each stencil point. compute_finalize_fn_closure: Closure finalizing the output value.

Parameters

• ​shape_element_type (DType): The element dtype of the shape.
• ​input_shape_element_type (DType): The element dtype of the input shape.
• ​rank (Int): Input and output domain rank.
• ​stencil_rank (Int): Rank of stencil subdomain slice.
• ​stencil_axis (IndexList[stencil_rank, element_type=element_type]): Stencil subdomain axes.
• ​simd_width (Int): The SIMD vector width to use.
• ​dtype (DType): The input and output data dtype.
• ​map_fn (def(IndexList[stencil_rank, element_type=*?]) -> Tuple[IndexList[stencil_rank], IndexList[stencil_rank]]): A function that a point in the output domain to the input co-domain.
• ​map_strides (def(dim: Int) -> Int): A function that returns the stride for the dim.
• ​load_fn (def[simd_width: Int, dtype: DType](IndexList[rank, element_type=*?]) -> SIMD[dtype, simd_width]): A function that loads a vector of simd_width from input.
• ​compute_init_fn (def[simd_width: Int]() -> SIMD[dtype, simd_width]): A function that initializes vector compute over the stencil.
• ​compute_fn (def[simd_width: SIMDLength](IndexList[rank, element_type=*?], SIMD[dtype, simd_width], SIMD[dtype, simd_width]) -> SIMD[dtype, simd_width]): A function the process the value computed for each point in the stencil.
• ​compute_finalize_fn (def[simd_width: SIMDLength](IndexList[rank, element_type=*?], SIMD[dtype, simd_width]) -> None): A function that finalizes the computation of a point in the output domain given a stencil.

stencil_gpu

comptime stencil_gpu[shape_element_type: DType, input_shape_element_type: DType, //, rank: Int, stencil_rank: Int, stencil_axis: IndexList[stencil_rank, element_type=element_type], simd_width: Int, dtype: DType, MapFnType: def(IndexList[stencil_rank, element_type=*?]) -> Tuple[IndexList[stencil_rank], IndexList[stencil_rank]] & ImplicitlyCopyable & RegisterPassable, MapStridesType: def(dim: Int) -> Int & ImplicitlyCopyable & RegisterPassable, LoadFnType: def[simd_width: Int, dtype: DType](IndexList[rank, element_type=*?]) -> SIMD[dtype, simd_width] & ImplicitlyCopyable & RegisterPassable, ComputeInitFnType: def[simd_width: Int]() -> SIMD[dtype, simd_width] & ImplicitlyCopyable & RegisterPassable, ComputeFnType: def[simd_width: SIMDLength](IndexList[rank, element_type=*?], SIMD[dtype, simd_width], SIMD[dtype, simd_width]) -> SIMD[dtype, simd_width] & ImplicitlyCopyable & RegisterPassable, ComputeFinalizeFnType: def[simd_width: SIMDLength](IndexList[rank, element_type=*?], SIMD[dtype, simd_width]) -> None & ImplicitlyCopyable & RegisterPassable] where (eq MapFnType.stencil_rank, stencil_rank) where (eq LoadFnType.rank, rank) where (eq ComputeInitFnType.dtype, dtype) where (eq ComputeFnType.rank, rank) where (eq ComputeFnType.dtype, dtype) where (eq ComputeFinalizeFnType.rank, rank) where (eq ComputeFinalizeFnType.dtype, dtype) = fn_literal

(Naive implementation) Computes stencil operation in parallel on GPU.

Args: ctx: The DeviceContext to use for GPU execution. shape: The shape of the output buffer. input_shape: The shape of the input buffer. map_func: Closure mapping output points to input co-domain bounds. map_strides_func: Closure returning the stride for a given dimension. load_func: Closure loading a SIMD vector from input. compute_init_func: Closure initializing the stencil accumulator. compute_func: Closure processing each stencil point. compute_finalize_func: Closure finalizing the output value.

Raises: If the GPU kernel launch fails.

Parameters

• ​shape_element_type (DType): The element dtype of the shape.
• ​input_shape_element_type (DType): The element dtype of the input shape.
• ​rank (Int): Input and output domain rank.
• ​stencil_rank (Int): Rank of stencil subdomain slice.
• ​stencil_axis (IndexList[stencil_rank, element_type=element_type]): Stencil subdomain axes.
• ​simd_width (Int): The SIMD vector width to use.
• ​dtype (DType): The input and output data dtype.
• ​MapFnType (def(IndexList[stencil_rank, element_type=*?]) -> Tuple[IndexList[stencil_rank], IndexList[stencil_rank]] & ImplicitlyCopyable & RegisterPassable): A closure maps a point in the output domain to input co-domain bounds.
• ​MapStridesType (def(dim: Int) -> Int & ImplicitlyCopyable & RegisterPassable): A closure returns the stride for each dimension.
• ​LoadFnType (def[simd_width: Int, dtype: DType](IndexList[rank, element_type=*?]) -> SIMD[dtype, simd_width] & ImplicitlyCopyable & RegisterPassable): A closure loads a SIMD vector from input.
• ​ComputeInitFnType (def[simd_width: Int]() -> SIMD[dtype, simd_width] & ImplicitlyCopyable & RegisterPassable): A closure initializes the stencil accumulator.
• ​ComputeFnType (def[simd_width: SIMDLength](IndexList[rank, element_type=*?], SIMD[dtype, simd_width], SIMD[dtype, simd_width]) -> SIMD[dtype, simd_width] & ImplicitlyCopyable & RegisterPassable): A closure processes the value computed for each stencil point.
• ​ComputeFinalizeFnType (def[simd_width: SIMDLength](IndexList[rank, element_type=*?], SIMD[dtype, simd_width]) -> None & ImplicitlyCopyable & RegisterPassable): A closure finalizes the output value from the stencil result.

Functions

• ​dual_elementwise: Executes two elementwise functions over their respective shapes in a single GPU kernel launch. Each thread processes elements from both shapes, fusing two independent elementwise passes into one.
• ​elementwise: 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.

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit