For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo function
multimem_ld_reduce
def multimem_ld_reduce[dtype: DType, *, count: Int, reduction: ReduceOp, scope: Scope, consistency: Consistency, accum_type: DType = get_accum_type[dtype](), output_width: Int = Int(1)](addr: Pointer[Scalar[dtype], address_space=AddressSpace.GLOBAL]) -> StaticTuple[SIMD[dtype, output_width], count]
Performs a vectorized load-reduce operation using NVIDIA's multimem feature.
This function loads multiple values from global memory and performs a reduction operation across them in a single instruction. It utilizes NVIDIA's multimem feature available on SM90+ GPUs for improved performance.
Constraints:
- Only supported on SM90+ GPUs.
- Total bit width (count * output_width * size_of[dtype] * 8) must be 32, 64, or 128 bits.
- Type must be a floating point type.
- float64 requires count=1 (no .vec qualifier allowed).
Parameters:
- dtype (
DType): Data dtype for the operation (must be a floating point type). - count (
Int): Vector size for PTX (corresponds to .v2, .v4, .v8 qualifiers, or no .v for scalar). - reduction (
ReduceOp): Type of reduction operation to perform. - scope (
Scope): Memory scope for the operation. - consistency (
Consistency): Memory consistency model to use. - accum_type (
DType): Data dtype used for accumulation. Defaults to a wider dtype than input (e.g. float32 for float16 inputs) to maintain precision during reduction. - output_width (
Int): Number of elements packed into a single output register (e.g. bf16x2).
Args:
- addr (
Pointer[Scalar[dtype], address_space=AddressSpace.GLOBAL]): Pointer to global memory where data will be loaded from.
Returns:
StaticTuple[SIMD[dtype, output_width], count]: A StaticTuple containing 'count' SIMD vectors of width 'output_width'
holding the results of the load-reduce operation.
def multimem_ld_reduce[dtype: DType, *, simd_width: Int, reduction: ReduceOp, scope: Scope, consistency: Consistency, accum_type: DType = get_accum_type[dtype]()](addr: Pointer[Scalar[dtype], address_space=AddressSpace.GLOBAL]) -> SIMD[dtype, simd_width]
Simplified multimem_ld_reduce that automatically calculates optimal packing.
This wrapper automatically determines the optimal output_width and count parameters based on the requested simd_width and data type, using 32-bit word packing for efficiency.
Constraints:
- Only supported on SM90+ GPUs.
- simd_width must be 1, 2, 4, or 8.
- Total bit width (count * output_width * size_of[dtype] * 8) must be 32, 64, or 128 bits.
- Type must be a floating point type.
- float64 requires count=1 (no .vec qualifier allowed).
Parameters:
- dtype (
DType): Data type of the elements to load (must be a floating point type). - simd_width (
Int): Total number of elements to process (must be 1, 2, 4, or 8). - reduction (
ReduceOp): Type of reduction operation to perform. - scope (
Scope): Memory scope for the operation. - consistency (
Consistency): Memory consistency model to use. - accum_type (
DType): Data type used for accumulation.
Args:
- addr (
Pointer[Scalar[dtype], address_space=AddressSpace.GLOBAL]): Pointer to global memory where data will be loaded from.
Returns:
SIMD[dtype, simd_width]: A SIMD vector containing simd_width elements with the reduction results.