# Copyright 2024 The JAX Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== from __future__ import annotations from collections.abc import Callable, Sequence import contextlib import dataclasses import enum import functools import math from typing import cast, Any, ClassVar, Literal from jax._src.lib import mosaic_gpu_dialect as mgpu_dialect from jax._src.lib import version as jaxlib_version from jaxlib.mlir import ir from jaxlib.mlir.dialects import _gpu_ops_gen from jaxlib.mlir.dialects import arith from jaxlib.mlir.dialects import builtin from jaxlib.mlir.dialects import func from jaxlib.mlir.dialects import gpu from jaxlib.mlir.dialects import llvm from jaxlib.mlir.dialects import memref from jaxlib.mlir.dialects import nvvm import numpy as np from . import fragmented_array as fa from . import profiler from . import utils TMA_DESCRIPTOR_BYTES = 128 TMA_DESCRIPTOR_ALIGNMENT = 64 # Matches `mosaic::gpu::kExpectedHbmAlignment`. HBM_ALIGNMENT = 256 TMAReductionOp = Literal[ "add", "min", "max", "inc", "dec", "and", "or", "xor", "umin", "umax", "smin", "smax", ] # Fixed size of the collective metadata structure in XLA. # Stores rank, param_to_peers_ptrs and multicast_buffer_ptr. # Plus alignment to 16 bytes. COLLECTIVE_METADATA_SIZE = 4 # Attribute used to merk the module which uses collective metadata. COLLECTIVE_ATTR = "mosaic_gpu.collective_metadata_used" # Attribute used to cache the kernel argument which corresponds to a given # reference during remote_ref lowering. KERNEL_ARG_ID_ATTR = "mosaic_gpu.from_kernel_arg_idx" # Attribute used to mark the first creation of the GMEM kernel arguments. ORIGINAL_KERNEL_ARG_ATTR = "mosaic_gpu.original_kernel_arg" # Attribute used to identify an operation used to load the current device id. # Needed for _recompute_peer_device_id in TMA descriptor construction. DEVICE_ID_ATTR = "mosaic_gpu.device_id_load" # Attribute used to mark that a kernel requires multicast support USES_MULTIMEM_ATTR = "mosaic_gpu.multimem_used" # Module attribute used to identify which kernel arguments are used with # multimem or used accross several processes. # Parameter is used as a bitset where the i-th bit is set if the i-th argument # must be allocated in symmetric memory space. SYMMETRIC_MEMORY_ARGS_ATTR = "mosaic_gpu.symmetric_memory_args" def uses_collective_metadata(module): return COLLECTIVE_ATTR in module.operation.attributes.keys() def _reduction_op_to_ptx(reduction_op: TMAReductionOp) -> str: # convert [s|u]min|max to min|max return reduction_op[-3:] c = utils.c # This is too common to fully qualify. class GlobalBroadcast: pass GLOBAL_BROADCAST = GlobalBroadcast() _DYNAMIC_PEER_TABLE = object() class CopyPartition: PARTITIONED: ClassVar[type[_Partitioned]] REPLICATED: ClassVar[_Replicated] @dataclasses.dataclass(frozen=True) class _Partitioned(CopyPartition): axis: int @dataclasses.dataclass(frozen=True) class _Replicated(CopyPartition): pass CopyPartition.PARTITIONED = _Partitioned CopyPartition.REPLICATED = _Replicated() @dataclasses.dataclass(frozen=True) class MemRefTransform: def apply(self, ref: ir.Value) -> ir.Value: raise NotImplementedError("Subclasses should override this method") def transform_index(self, idx: Sequence[ir.Value]) -> tuple[ir.Value, ...]: raise NotImplementedError("Subclasses should override this method") def transform_shape(self, shape: Sequence[int]) -> tuple[int, ...]: raise NotImplementedError("Subclasses should override this method") def transform_gmem_shape(self, shape: Sequence[int]) -> tuple[int, ...]: """Applies the shape transformation to the given GMEM shape. This function is intended to mirror the behavior of the `apply` method on GMEM shapes. """ raise NotImplementedError("Subclasses should override this method") def transform_strides(self, strides: Sequence[int]) -> tuple[int, ...]: raise NotImplementedError("Subclasses should override this method") def batch(self, leading_rank: int) -> MemRefTransform: """Returns a transform that accepts a ref with the extra `leading_rank` dims. The returned transform should leave the leading dimensions unchanged and only apply to the suffix of the shape. """ raise NotImplementedError("Subclasses should override this method") def to_attr(self) -> ir.Attribute: raise NotImplementedError( f"Subclasses should override this method: {type(self)}" ) class Rounding(enum.Enum): UP = enum.auto() DOWN = enum.auto() @dataclasses.dataclass(frozen=True) class TileTransform(MemRefTransform): """Tiles a suffix of memref dimensions. For example, given a memref of shape (5, 128, 128) and a tiling of (64, 32), the shape of the result will be (5, 2, 4, 64, 32). The shape always ends with the tile shape, and the size of tiled dimensions is divided by the tile size. This is especially useful for swizzled WGMMA, which expect tiled layouts in shared memory. """ tiling: tuple[int, ...] rounding: Rounding | None = None def transform_gmem_shape(self, shape: Sequence[int]) -> tuple[int, ...]: untiled_rank = len(shape) tiling_rank = len(self.tiling) tiled_rank = untiled_rank + tiling_rank shape = list(shape) for t, d in zip(self.tiling[::-1], range(untiled_rank)[::-1]): s = shape[d] if s > t: if s % t: match self.rounding: case None: raise ValueError( f"When no rounding mode is specified, dimension {d} must have" f" size smaller or a multiple of its tiling {t}, but got {s}" ) case Rounding.UP: raise NotImplementedError case Rounding.DOWN: s = s // t * t case _: raise ValueError(f"Unknown rounding mode: {self.rounding}") else: t = s shape[d : d + 1] = (s // t, t) permutation = ( *range(untiled_rank - tiling_rank), *range(untiled_rank - tiling_rank, tiled_rank, 2), *range(untiled_rank - tiling_rank + 1, tiled_rank, 2), ) return TransposeTransform(permutation).transform_shape(shape) def apply(self, ref: ir.Value) -> ir.Value: untiled_rank = ir.MemRefType(ref.type).rank tiling_rank = len(self.tiling) tiled_rank = untiled_rank + tiling_rank for t, d in zip(self.tiling[::-1], range(untiled_rank)[::-1]): ref_shape = ir.MemRefType(ref.type).shape s = ref_shape[d] if s > t: if s % t: match self.rounding: case None: raise ValueError( f"When no rounding mode is specified, dimension {d} must have" f" size smaller or a multiple of its tiling {t}, but got {s}" ) case Rounding.UP: raise NotImplementedError case Rounding.DOWN: slices = [slice(None)] * d slices.append(slice(0, s // t * t)) ref = utils.memref_slice(ref, tuple(slices)) case _: raise ValueError(f"Unknown rounding mode: {self.rounding}") else: t = s ref = utils.memref_unfold(ref, d, (None, t)) permutation = ( *range(untiled_rank - tiling_rank), *range(untiled_rank - tiling_rank, tiled_rank, 2), *range(untiled_rank - tiling_rank + 1, tiled_rank, 2), ) return utils.memref_transpose(ref, permutation) def transform_index(self, idx: Sequence[ir.Value]) -> tuple[ir.Value, ...]: index = ir.IndexType.get() tiling_rank = len(self.tiling) return ( *idx[:-tiling_rank], *( arith.divui(i, c(t, index)) for i, t in zip(idx[-tiling_rank:], self.tiling) ), *( arith.remui(i, c(t, index)) for i, t in zip(idx[-tiling_rank:], self.tiling) ), ) def transform_shape(self, shape: Sequence[int]) -> tuple[int, ...]: # Note that this also checks that tiled dims are not squeezed. Their slice # size would be 1 if so. tiling_rank = len(self.tiling) if self.rounding is None: for size, tile_size in zip(shape[-tiling_rank:], self.tiling, strict=True): if size % tile_size: raise ValueError( f"Expected GMEM slice shape {shape} suffix to be a multiple of" f" tiling {self.tiling}.\nIf you're using padded async copies," " your slice might need to extend out of bounds of the GMEM" " buffer (OOB accesses will be skipped)." ) elif self.rounding != Rounding.DOWN: raise NotImplementedError(self.rounding) return ( *shape[:-tiling_rank], *(s // t for s, t in zip(shape[-tiling_rank:], self.tiling)), *self.tiling, ) def transform_strides(self, strides: Sequence[int]) -> tuple[int, ...]: tiling_rank = len(self.tiling) return ( *strides[:-tiling_rank], *(s * t for s, t in zip(strides[-tiling_rank:], self.tiling)), *strides[-tiling_rank:], ) def batch(self, leading_rank: int) -> MemRefTransform: return self def to_attr(self) -> ir.Attribute: return mgpu_dialect.TileTransformAttr.get(self.tiling) @dataclasses.dataclass(frozen=True) class TransposeTransform(MemRefTransform): """Transposes memref dimensions.""" permutation: tuple[int, ...] def __post_init__(self): if len(self.permutation) != len(set(self.permutation)): raise ValueError("All elements of `permutation` must be unique") def apply(self, ref: ir.Value) -> ir.Value: return utils.memref_transpose(ref, self.permutation) def transform_gmem_shape(self, shape: Sequence[int]) -> tuple[int, ...]: return self.transform_shape(shape) def transform_index(self, idx: Sequence[ir.Value]) -> tuple[ir.Value, ...]: return tuple(idx[p] for p in self.permutation) def transform_shape(self, shape: Sequence[int]) -> tuple[int, ...]: return tuple(shape[p] for p in self.permutation) def transform_strides(self, strides: Sequence[int]) -> tuple[int, ...]: return tuple(strides[p] for p in self.permutation) def batch(self, leading_rank: int) -> MemRefTransform: return TransposeTransform( (*range(leading_rank), *(d + leading_rank for d in self.permutation)) ) @dataclasses.dataclass(frozen=True) class CollapseLeadingIndicesTransform(MemRefTransform): """Collapses leading indices into one.""" strides: tuple[int, ...] @functools.cached_property def common_stride(self) -> int: return math.gcd(*self.strides) def apply(self, ref: ir.Value) -> ir.Value: ref_ty = ir.MemRefType(ref.type) strides, offset = ref_ty.get_strides_and_offset() if offset == ir.ShapedType.get_dynamic_stride_or_offset(): raise NotImplementedError("Dynamic offsets are not supported") max_bound = sum( (d - 1) * s // self.common_stride for d, s in zip( ref_ty.shape[: len(self.strides)], strides[: len(self.strides)] ) ) + 1 new_shape = [max_bound, *ref_ty.shape[len(self.strides):]] new_strides = [self.common_stride, *strides[len(self.strides):]] new_layout = ir.StridedLayoutAttr.get(offset, new_strides) new_ref_ty = ir.MemRefType.get( new_shape, ref_ty.element_type, new_layout, ref_ty.memory_space ) return memref.reinterpret_cast( new_ref_ty, ref, [], [], [], static_offsets=[offset], static_sizes=new_shape, static_strides=new_strides, ) def transform_gmem_shape(self, shape: Sequence[int]) -> tuple[int, ...]: raise NotImplementedError # Unused def transform_index(self, idx: Sequence[ir.Value]) -> tuple[ir.Value, ...]: index = ir.IndexType.get() flat_idx = c(0, index) for i, s in zip(idx[:len(self.strides)], self.strides): flat_idx = arith.addi( flat_idx, arith.muli(i, c(s // self.common_stride, index)) ) return (flat_idx, *idx[len(self.strides):]) def transform_shape(self, shape: Sequence[int]) -> tuple[int, ...]: if any(s != 1 for s in shape[:len(self.strides)]): raise ValueError("Expected leading indices to be squeezed") return (1, *shape[len(self.strides):]) def batch(self, leading_rank: int) -> MemRefTransform: raise NotImplementedError # Unused @dataclasses.dataclass(frozen=True) class DropUnitDimsTransform(MemRefTransform): """Drops unit dimensions at the given positions.""" unit_dims: tuple[int, ...] def apply(self, ref: ir.Value) -> ir.Value: ref_ty = ir.MemRefType(ref.type) strides, offset = ref_ty.get_strides_and_offset() rank = len(ref_ty.shape) # This will assert that all the relevant dimensions are indeed trivial. new_shape = self.transform_shape(ref_ty.shape) new_strides = [strides[i] for i in range(rank) if i not in self.unit_dims] # Not propagating offset, as it is folded into the ref by `memref_ptr`. new_layout = ir.StridedLayoutAttr.get(offset, new_strides) new_ref_ty = ir.MemRefType.get( new_shape, ref_ty.element_type, new_layout, ref_ty.memory_space ) if offset == ir.ShapedType.get_dynamic_stride_or_offset(): _, dyn_offset, *_ = memref.extract_strided_metadata(ref) # pyrefly: ignore[not-iterable] offsets = [dyn_offset] else: offsets = [] return memref.reinterpret_cast( new_ref_ty, ref, offsets, [], [], static_offsets=[offset], static_sizes=new_shape, static_strides=new_strides, ) def transform_gmem_shape(self, shape: Sequence[int]) -> tuple[int, ...]: return self.transform_shape(shape) def transform_index(self, idx: Sequence[ir.Value]) -> tuple[ir.Value, ...]: return tuple(v for i, v in enumerate(idx) if i not in self.unit_dims) def transform_shape(self, shape: Sequence[int]) -> tuple[int, ...]: for d in self.unit_dims: if shape[d] != 1: raise ValueError( f"Expected dimension {d} to have size 1 in shape {shape}, but got" f" {shape[d]}" ) return tuple(v for i, v in enumerate(shape) if i not in self.unit_dims) def batch(self, leading_rank: int) -> MemRefTransform: raise NotImplementedError # Unused OnDeviceProfiler = profiler.OnDeviceProfiler MOSAIC_GPU_SMEM_ALLOC_ATTR = "mosaic_gpu_smem_alloc" class Scratch: """Manages ops handling the GMEM scratch that contains the TMA descriptors. TMA descriptors are created on the host and then copied to GMEM. So there needs to be some code on the host to allocate and initialize the TMA descriptors. However, we only know what descriptors we need after we have lowered the entire kernel. This class helps manage everything needed to correctly allocate and initialize the scratch. To help reconcile the needs of kernels that use the dialect lowering with those that use MGPU APIs directly, this class only creates the relevant ops lazily. Eager creation would make them appear dead before dialect lowering and MLIR's DCE would remove them. During the lowering, we collect information about how many bytes are needed and also how each descriptor should be initialized on the host. At the end of the lowering, the finalize_size() method should be called to add the necessary code on the host to allocate and initialize all descriptors. Here's how the IR looks after the initial ops are created for the first time: %1 = llvm.alloc_op {elem_type = !llvm.array<0 x i8>} -> !llvm.ptr %2 = llvm.load_op (%1) : (!llvm.ptr) -> !llvm.array<0 x i8> ... %3 = gpu.launch async ^bb0: %4 = builtin.unrealized_conversion_cast_op(%2) : (!llvm.array<256 x i8>) -> !llvm.ptr And here is an example of how the IR could look like after finalize_size() is called: %11 = llvm.alloc_op {elem_type = !llvm.array<256 x i8>} -> !llvm.ptr %22 = llvm.load_op (%11) : (!llvm.ptr) -> !llvm.array<256 x i8> ... # Ops inserted to initialize the tma descriptors on the host: ... %33 = llvm.getelementptr %11[0] : (!llvm.ptr) -> !llvm.ptr, i8 call @mosaic_gpu_init_tma_desc (%33, ...) ... %44 = llvm.getelementptr %11[128] : (!llvm.ptr) -> !llvm.ptr, i8 call @mosaic_gpu_init_tma_desc (%44, ...) ... %55 = gpu.launch async ^bb0: %66 = builtin.unrealized_conversion_cast_op(%22) : (!llvm.array<256 x i8>) -> !llvm.ptr """ def __init__(self, gpu_launch_op: _gpu_ops_gen.LaunchOp): self.next_offset: int = 0 self.host_init: list[Callable[[ir.Value], None]] = [] self.descriptor_offsets: list[int] = [] self.descriptor_prefetch: list[bool] = [] self._ops_created = False # Ideally, we would store the gpu.launch op directly. However, it gets # invalidated by passes like "canonicalize". Thus we store the module and # find the gpu.launch op from there when needed. op = gpu_launch_op while op.name != "builtin.module": assert op.parent is not None op = op.parent.opview assert op is not None self._module_op = cast(builtin.ModuleOp, op) def _find_first_op[T: ir.OpView]( self, op_type: type[T], block: ir.Block, tag_attribute_name: str | None = None ) -> T | None: op_name = getattr(op_type, "OPERATION_NAME", None) for op in block: if op.name == op_name and ( tag_attribute_name is None or tag_attribute_name in op.attributes ): return cast(T, op) for region in op.regions: for block in region: child_op = self._find_first_op(op_type, block, tag_attribute_name) if child_op is not None: return child_op return None def _create_ops(self): if self._ops_created: return self._ops_created = True gpu_launch_op = self._find_first_op(gpu.LaunchOp, self._module_op.body) assert gpu_launch_op is not None i8 = ir.IntegerType.get_signless(8) i64 = ir.IntegerType.get_signless(64) ptr_ty = llvm.PointerType.get() empty_arr_ty = llvm.ArrayType.get(i8, 0) with ir.InsertionPoint(gpu_launch_op): alloc_op = llvm.AllocaOp( ptr_ty, c(1, i64), empty_arr_ty, alignment=TMA_DESCRIPTOR_ALIGNMENT ) # Tag the alloc op with an attribute so that we can find it later. alloc_op.attributes[MOSAIC_GPU_SMEM_ALLOC_ATTR] = ir.UnitAttr.get() load_op = llvm.LoadOp(empty_arr_ty, alloc_op.result) with ir.InsertionPoint.at_block_begin(gpu_launch_op.body.blocks[0]): builtin.unrealized_conversion_cast([ptr_ty], [load_op.result]) def _find_alloc_load_and_device_ptr( self, ) -> tuple[llvm.AllocaOp, llvm.LoadOp, ir.OpResult[llvm.PointerType]]: if not self._ops_created: self._create_ops() alloc_op = self._find_first_op( llvm.AllocaOp, self._module_op.body, MOSAIC_GPU_SMEM_ALLOC_ATTR ) assert alloc_op is not None [alloc_user] = alloc_op.result.uses load_op = alloc_user.owner assert load_op.operation.name == "llvm.load" [load_op_user] = load_op.result.uses device_ptr = load_op_user.owner assert device_ptr.operation.name == "builtin.unrealized_conversion_cast" return alloc_op, load_op, device_ptr.result # pyrefly: ignore[bad-return] def device_ptr(self) -> ir.Value: _, _, device_ptr = self._find_alloc_load_and_device_ptr() return device_ptr def finalize_size(self): """ Allocates and initializes the host buffer. This needs to be done after lowering, i.e. after all TMA descriptors have been recorded. Only then we know what the scratch contains. """ if self.next_offset == 0: return alloc_op, load_op, device_ptr = self._find_alloc_load_and_device_ptr() i8 = ir.IntegerType.get_signless(8) ptr_ty = llvm.PointerType.get() with ir.InsertionPoint(load_op): gmem_scratch_bytes = self.next_offset scratch_arr_ty = llvm.ArrayType.get(i8, gmem_scratch_bytes) alloc_op.elem_type = ir.TypeAttr.get(scratch_arr_ty) load_op.result.set_type(scratch_arr_ty) for init_callback in self.host_init: init_callback(alloc_op.result) with ir.InsertionPoint.after(device_ptr.owner): predicate = utils.single_thread_predicate(utils.ThreadSubset.BLOCK) for offset, prefetch in zip(self.descriptor_offsets, self.descriptor_prefetch, strict=True): if not prefetch: continue desc_ptr = llvm.getelementptr( ptr_ty, device_ptr, [], [offset], i8, llvm.GEPNoWrapFlags.none ) utils.prefetch_tensormap(desc_ptr, predicate=predicate) class _DefaultPredicate: pass def _find_kernel_argument_for_gmem_ref(gmem_ref: ir.Value) -> ir.Value: """Returns the kernel argument value for a given gmem_ref. The kernel argument is expected to be an unrealized conversion cast. This function will recursively go up block arguments in case of nested blocks. """ if not isinstance(gmem_ref.type, ir.MemRefType): raise ValueError(f"Expected {gmem_ref} to have a memref type.") while isinstance(gmem_ref, ir.BlockArgument): gmem_ref = gmem_ref.owner.owner.operands[gmem_ref.arg_number] if ORIGINAL_KERNEL_ARG_ATTR not in gmem_ref.owner.attributes: # pyrefly: ignore[missing-attribute] raise NotImplementedError( f"Expected {gmem_ref.owner} to be a GMEM kernel argument." ) return gmem_ref def _is_tma_reduction_op_supported( reduction_op: TMAReductionOp | None, dtype: ir.Type, ) -> bool: """Returns whether the given TMA reduction op supports the given dtype. This function essentially implements the table at: https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-reduce-async-bulk-tensor with the following differences: - For `add` reductions, we also support int64, treating it as uint64. - For `and`, `or`, and `xor` reductions, we support signed integer types. - For `inc` and `dec` reductions, we support both signed and unsigned i32 treating both as unsigned. """ i32 = ir.IntegerType.get_signless(32) i64 = ir.IntegerType.get_signless(64) f16 = ir.F16Type.get() f32 = ir.F32Type.get() bf16 = ir.BF16Type.get() match reduction_op: case None: return True case "add": return dtype in (f16, f32, bf16, i32, i64) case "max" | "min": return dtype in (f16, bf16) case "umax" | "umin" | "smax" | "smin": return dtype in (i32, i64) case "inc" | "dec": return dtype == i32 case "and" | "or" | "xor": return dtype in (i32, i64) def _tma_dma_type( element_type: ir.Type, reduction_op: TMAReductionOp | None, ) -> int: """Returns the TMA DMA type for the given element type and signedness.""" if isinstance(element_type, ir.IntegerType): bitwidth = utils.bitwidth_impl(element_type) if bitwidth == 2: tma_dtype = 8 elif bitwidth == 4: tma_dtype = 0 elif bitwidth == 8: tma_dtype = 1 elif bitwidth == 16: tma_dtype = 2 elif bitwidth == 32: tma_dtype = 9 if reduction_op in ("smin", "smax") else 3 elif bitwidth == 64: tma_dtype = 10 if reduction_op in ("smin", "smax") else 4 else: raise ValueError(f"Unsupported integer bitwidth: {bitwidth}") elif isinstance(element_type, ir.F16Type): tma_dtype = 5 elif isinstance(element_type, ir.F32Type): tma_dtype = 6 elif isinstance(element_type, ir.BF16Type): tma_dtype = 7 # We treat narrow floats as integers elif isinstance(element_type, ir.Float8E5M2Type): tma_dtype = 1 elif isinstance(element_type, ir.Float8E4M3FNType): tma_dtype = 1 elif isinstance(element_type, ir.Float8E8M0FNUType): tma_dtype = 1 elif isinstance(element_type, ir.Float4E2M1FNType): tma_dtype = 0 else: raise ValueError(f"unsupported TMA dtype {element_type}") return tma_dtype class AsyncCopyImplementation(enum.Enum): TMA = enum.auto() CP_ASYNC = enum.auto() def get_collective_metadata_size(num_params: int, num_peers: int) -> int: """Returns the size of the collective metadata buffer for the given number of parameters and peers.""" return ( # Stores the collective metadata structure. COLLECTIVE_METADATA_SIZE # For each peer we need to store a pointer to each parameter. + num_peers * num_params # For each parameter we need to store a pointer to the multimem address. + num_params ) class OOBFillMode(enum.IntEnum): UNDEFINED = 0 PROMISE_IN_BOUNDS = 1 ZEROS = 2 @dataclasses.dataclass() class LaunchContext: module: ir.Module scratch: Scratch cluster_size: tuple[int, int, int] buffers: ir.Value profiler: OnDeviceProfiler | None = None num_peers: int = 0 num_params: int = 0 tma_descriptors: dict[Any, ir.Value] = dataclasses.field( default_factory=dict, init=False ) is_device_collective: bool = False multihost_kernel: bool = False @contextlib.contextmanager def named_region(self, *args, **kwargs): if self.profiler is not None: with self.profiler.record(*args, **kwargs): yield else: yield def _collective_metadata(self, idx: int) -> ir.Value | None: if self.num_peers <= 1: return None host_block = self.buffers.owner assert isinstance(host_block, ir.Block) with ir.InsertionPoint.at_block_begin(host_block): ptr_ty = llvm.PointerType.get() buffer_ptr = utils.getelementptr( self.buffers, [self.num_params + idx], ptr_ty ) metadata_ptr = llvm.load(ptr_ty, buffer_ptr) metadata_ty = ir.MemRefType.get( (get_collective_metadata_size(self.num_params, self.num_peers),), ir.IntegerType.get_signless(64), ) return utils.ptr_as_memref(metadata_ptr, metadata_ty) @functools.cached_property def device_collective_metadata(self) -> ir.Value | None: return self._collective_metadata(0) @functools.cached_property def host_collective_metadata(self) -> ir.Value | None: return self._collective_metadata(1) def _alloc_scratch( self, size: int, alignment: int | None = None, host_init: Callable[[ir.Value], None] = lambda _: None, prefetch: bool = True, ) -> ir.Value: """Allocates a GMEM scratch buffer. The buffer is initialized on the host and then copied to GMEM before the kernel launch. """ i8 = ir.IntegerType.get_signless(8) ptr_ty = llvm.PointerType.get() if alignment is None: alignment = size if self.scratch.next_offset % alignment: raise NotImplementedError # TODO(apaszke): Pad to match alignment alloc_base = self.scratch.next_offset self.scratch.next_offset += size def host_init_wrapped(host_ptr): host_init( llvm.getelementptr( ptr_ty, host_ptr, [], [alloc_base], i8, llvm.GEPNoWrapFlags.none ), ) self.scratch.host_init.append(host_init_wrapped) self.scratch.descriptor_offsets.append(alloc_base) self.scratch.descriptor_prefetch.append(prefetch) # with ir.InsertionPoint(self.gmem_scratch_ptr.owner): # There is no way to create an insertion point after an operation... gep = llvm.GEPOp( ptr_ty, self.scratch.device_ptr(), [], [alloc_base], i8, llvm.GEPNoWrapFlags.none ) gep.move_after(self.scratch.device_ptr().owner) # pyrefly: ignore[bad-argument-type] return gep.result def _recompute_peer_id( self, peer_id: ir.Value, fuel=8, dry_run: bool = False, ) -> ir.Value: # In dry_run mode, this function purely verifies whether the expression # can be reproduced on the host without mutating the MLIR module (incl. # emitting operations) and returns peer_id unchanged. if fuel == 0: raise ReplicationError( "gmem_peer_id computation is too complicated to recompute on the host" ) if isinstance(peer_id, ir.BlockArgument): raise ReplicationError("Can't recompute a value that's a block argument") op: Any = peer_id.owner # We accept all arith ops if op.OPERATION_NAME.startswith("arith."): if DEVICE_ID_ATTR in op.attributes: return peer_id if dry_run else self.device_id(on_host=True) new_operands = [ self._recompute_peer_id(x, fuel - 1, dry_run) for x in op.operands ] if dry_run: return peer_id result_types = [r.type for r in op.results] new_attributes = {na: op.attributes[na] for na in op.attributes} new_op = ir.Operation.create( op.OPERATION_NAME, result_types, new_operands, new_attributes ) assert len(new_op.results) == 1 return new_op.result # nvshmem_my_pe queries the device id of the current process and works on # both the host and the device. if ( isinstance(op, llvm.CallOp) and op.callee is not None and op.callee.value == "nvshmem_my_pe" ): if dry_run: return peer_id i32 = ir.IntegerType.get_signless(32) return cast(ir.Value, llvm.call(i32, [], [], [], callee="nvshmem_my_pe")) raise ReplicationError( f"Unrecognized op can't be recomputed on the host: {op}" ) def _flag_multimem_usage(self): self.module.operation.attributes[USES_MULTIMEM_ATTR] = ir.UnitAttr.get() def _get_tma_desc( self, gmem_ref: ir.Value, gmem_transform: tuple[MemRefTransform, ...], gmem_peer_id: int | ir.Value | GlobalBroadcast | None, transformed_slice_shape: tuple[int, ...], swizzle: int | None, reduction_op: TMAReductionOp | None, *, window_start_bounds: tuple[tuple[int, int], ...] | None = None, window_strides: tuple[int, ...] | None = None, ): gmem_ref = _find_kernel_argument_for_gmem_ref(gmem_ref) tma_dtype = _tma_dma_type(ir.MemRefType(gmem_ref.type).element_type, reduction_op) i8 = ir.IntegerType.get_signless(8) i32 = ir.IntegerType.get_signless(32) i64 = ir.IntegerType.get_signless(64) ptr_ty = llvm.PointerType.get() is_dynamic_peer = False if isinstance(gmem_peer_id, ir.Value): if gmem_peer_id.type != i32: raise TypeError(f"Expected peer id to be an i32, got {gmem_peer_id.type}") try: self._recompute_peer_id(gmem_peer_id, fuel=16, dry_run=True) except ReplicationError: is_dynamic_peer = True # Without a peer count (e.g. NVSHMEM) we can't size the descriptor table. if is_dynamic_peer and self.num_peers == 0: raise NotImplementedError( "Peer ids that can't be recomputed on the host are unsupported" " when the number of peers is unknown at compile time." ) # One host-side descriptor per peer is too expensive on e.g. NVL72. if is_dynamic_peer and self.num_peers > 8: raise NotImplementedError( "Dynamic peer async_copy currently supports at most 8 peers, but" f" got {self.num_peers} peers." ) # Using ir.Values in cache keys is a little sketchy, but I think it should # be fine. Having it in the key will keep it alive, and if comparison and # hashing is by identity then it should work out. tma_desc_key = ( gmem_ref, transformed_slice_shape, window_start_bounds, window_strides, swizzle, gmem_transform, _DYNAMIC_PEER_TABLE if is_dynamic_peer else gmem_peer_id, tma_dtype, ) if (tma_desc := self.tma_descriptors.get(tma_desc_key, None)) is None: def init_tma_desc(host_ptr: ir.Value): ref = gmem_ref for t in gmem_transform: ref = t.apply(ref) ref_ty = ir.MemRefType(ref.type) # TODO(apaszke): Use utils.memref_ptr to compute base_ptr strides, _ = ref_ty.get_strides_and_offset() if strides[-1] != 1: raise ValueError( "TMA requires the stride of the last dimension after" " transforming the GMEM reference to be 1, but it is" f" {strides[-1]}." ) # pyrefly: ignore[not-iterable] _, offset, *sizes_and_strides = memref.extract_strided_metadata(ref) aligned_ptr_idx = memref.extract_aligned_pointer_as_index(ref) as_i64 = lambda i: arith.index_cast(i64, i) alloc_ptr = llvm.inttoptr(ptr_ty, as_i64(aligned_ptr_idx)) llvm_dyn = -2147483648 # TODO(apaszke): Improve the MLIR bindings... base_ptr = llvm.getelementptr( ptr_ty, alloc_ptr, [as_i64(offset)], [llvm_dyn], ref_ty.element_type, llvm.GEPNoWrapFlags.none, ) rank = ref_ty.rank assert rank * 2 == len(sizes_and_strides) swizzle_arg = ( mgpu_dialect.SwizzlingMode.kNoSwizzle if swizzle is None else swizzle ) # TODO(apaszke): Better verification (e.g. slice is non-zero) # TODO(apaszke): We always know strides statically. dtype_or_bitwidth = c(tma_dtype, i64) def get_peer_base_ptr(peer_id: ir.Value) -> ir.Value: return utils.memref_ptr(self.to_remote(ref, peer_id, on_host=True)) def init_single_tma_desc(desc_ptr: ir.Value, target_base_ptr: ir.Value): args: list[ir.Value] = [ desc_ptr, target_base_ptr, dtype_or_bitwidth, c(rank, i64), utils.pack_array([as_i64(i) for i in sizes_and_strides[:rank]]), utils.pack_array([as_i64(i) for i in sizes_and_strides[rank:]]), c(swizzle_arg, i64), ] if window_start_bounds is not None: assert window_strides is not None num_rows, num_channels = transformed_slice_shape lower, upper = zip(*window_start_bounds) func.call( [], "mosaic_gpu_init_tma_im2col_desc", [ *args, utils.pack_array([c(v, i32) for v in lower]), utils.pack_array([c(v, i32) for v in upper]), c(num_channels, i64), c(num_rows, i64), utils.pack_array([c(v, i32) for v in window_strides]), ], ) else: func.call( [], "mosaic_gpu_init_tma_desc", [ *args, utils.pack_array( [c(v, i64) for v in transformed_slice_shape] ), ], ) if is_dynamic_peer: for p in range(self.num_peers): p_host_desc_ptr = utils.getelementptr( host_ptr, [p * TMA_DESCRIPTOR_BYTES], i8 ) init_single_tma_desc(p_host_desc_ptr, get_peer_base_ptr(c(p, i32))) elif isinstance(gmem_peer_id, GlobalBroadcast): multimem_ref = self.to_remote_multicast(ref, on_host=True) init_single_tma_desc(host_ptr, utils.memref_ptr(multimem_ref.ref)) elif gmem_peer_id is None: init_single_tma_desc(host_ptr, base_ptr) else: peer_id = ( self._recompute_peer_id(gmem_peer_id, fuel=16) if isinstance(gmem_peer_id, ir.Value) else c(gmem_peer_id, i32) ) init_single_tma_desc(host_ptr, get_peer_base_ptr(peer_id)) total_bytes = TMA_DESCRIPTOR_BYTES * (self.num_peers if is_dynamic_peer else 1) tma_desc = self._alloc_scratch( total_bytes, alignment=TMA_DESCRIPTOR_ALIGNMENT, host_init=init_tma_desc, prefetch=not is_dynamic_peer, ) self.tma_descriptors[tma_desc_key] = tma_desc if is_dynamic_peer: assert isinstance(gmem_peer_id, ir.Value) byte_offset = arith.muli(gmem_peer_id, c(TMA_DESCRIPTOR_BYTES, i32)) return utils.getelementptr(tma_desc, [byte_offset], i8) return tma_desc def _prepare_async_copy( self, gmem_ref: ir.Value, gmem_slice: Any, gmem_transform: tuple[MemRefTransform, ...], collective: tuple[gpu.Dimension, ...], leader_tracked: CopyPartition | None, implementation: AsyncCopyImplementation, ): """Performs setup common to TMA and CP_ASYNC implementations.""" index = ir.IndexType.get() gmem_ref_ty = ir.MemRefType(gmem_ref.type) gmem_strides, _ = gmem_ref_ty.get_strides_and_offset() if gmem_strides != utils.get_contiguous_strides(gmem_ref_ty.shape): raise NotImplementedError( "Async copies assume the GMEM reference is contiguous" ) # Look for and verify gather indices in gmem_slice. is_gathered_dim = [isinstance(s, fa.FragmentedArray) for s in gmem_slice] gather_indices: fa.FragmentedArray | None = None if any(is_gathered_dim): if is_gathered_dim != [True, False]: raise NotImplementedError( "Gathers/scatters only supported along the first dimension of 2D" " arrays" ) gather_indices = gmem_slice[0] if not isinstance(gather_indices, fa.FragmentedArray): raise ValueError("Gather/scatter indices must be a FragmentedArray") if len(gather_indices.shape) != 1: raise ValueError("Gather/scatter indices must be 1D") idx_dtype = gather_indices.mlir_dtype if ( not isinstance(idx_dtype, ir.IntegerType) or utils.bitwidth(idx_dtype) > 32 ): raise ValueError("Gather/scatter indices must be integers that are at most 32-bit wide") if gather_indices.is_signed: raise ValueError("Gather/scatter indices must be unsigned") gmem_slice = (slice(None), *gmem_slice[1:]) # Analyze the slice (taking gathers into account). base_indices, slice_shape, is_squeezed = utils.parse_indices( gmem_slice, ir.MemRefType(gmem_ref.type).shape, # NOTE: TMA supports OOB indices, so we skip the check. check_oob=implementation != AsyncCopyImplementation.TMA, ) if gather_indices is not None: slice_shape = [gather_indices.shape[0], *slice_shape[1:]] del gmem_slice # Use slice_shape, base_indices and is_squeezed from now on! dyn_base_indices = tuple( c(i, index) if not isinstance(i, ir.Value) else i for i in base_indices ) del base_indices # Use the dynamic indices from now on! # Deal with collective and partitioned loads. if collective: if implementation != AsyncCopyImplementation.TMA: raise ValueError("Only the TMA implementation supports collective copies") if gather_indices is not None: raise NotImplementedError("Collective copies with gather/scatter unsupported") if leader_tracked is not None: if not collective: # This implies non-gather TMA already. raise ValueError("Only collective loads can be leader-tracked") collective_size = math.prod(self.cluster_size[d] for d in collective) if collective_size > 1: if collective != (gpu.Dimension.x,) or collective_size != 2: raise NotImplementedError( "Leader-tracked loads only supported along the x dimension," " which must have size 2" ) if isinstance(leader_tracked, _Partitioned): partitioned = leader_tracked.axis # Increment partitioned by the number of preceding squeezed dimensions. partitioned = np.where( np.cumsum(~np.array(is_squeezed)) == partitioned+1)[0][0] # Partitioning happens on the logical slice we extract from GMEM, so we do # it before we apply transforms. if slice_shape[partitioned] % collective_size != 0: raise ValueError( f"The collective size ({collective_size}) must divide the slice" " shape along the partitioned dimension, but it has size" f" {slice_shape[partitioned]}" ) slice_shape[partitioned] //= collective_size dyn_base_indices = list(dyn_base_indices) dyn_base_indices[partitioned] = arith.addi( dyn_base_indices[partitioned], arith.muli( utils.cluster_idx(collective), c(slice_shape[partitioned], index), ), ) dyn_base_indices = tuple(dyn_base_indices) squeezed_dims = tuple( i for i, squeezed in enumerate(is_squeezed) if squeezed ) # Indexing is really slicing + squeezing, and user transforms are meant to # apply after that. However, we actually have to apply the indexing last # (it's fused into the TMA) and so we need to commute it with all the user # transforms. For slicing this is done using transform_index and # transform_shape. For squeezing we actually move all the squeezed dims to # the front, and then batch each transform, making it ignore the extra dims. if squeezed_dims and implementation != AsyncCopyImplementation.CP_ASYNC: sliced_dims = [i for i, squeezed in enumerate(is_squeezed) if not squeezed] gmem_transform = (TransposeTransform((*squeezed_dims, *sliced_dims)), *(t.batch(len(squeezed_dims)) for t in gmem_transform)) untransformed_slice_shape = tuple(slice_shape) slice_shape = untransformed_slice_shape for t in gmem_transform: dyn_base_indices = t.transform_index(dyn_base_indices) slice_shape = t.transform_shape(slice_shape) return ( list(slice_shape), untransformed_slice_shape, dyn_base_indices, squeezed_dims, gather_indices, gmem_transform, ) def _prepare_tma( self, gmem_ref: ir.Value, smem_ref: ir.Value | None, swizzle: int | None, slice_shape: list[int], dyn_base_indices: tuple[ir.Value, ...], gather_indices, squeezed_dims: tuple[int, ...], gmem_transform: tuple[MemRefTransform, ...], collective: tuple[gpu.Dimension, ...], leader_tracked: CopyPartition | None = None, ): """Finalizes setup specific to the TMA implementation of async_copy.""" index = ir.IndexType.get() # The function below is called only to verify the GMEM ref. The output # is meant to be ignored. _find_kernel_argument_for_gmem_ref(gmem_ref) gmem_ref_ty = ir.MemRefType(gmem_ref.type) gmem_strides, _ = gmem_ref_ty.get_strides_and_offset() # We don't need to do this for gather TMAs, because we'll unroll the # transfers ourselves anyway. num_squeezed_dims = len(squeezed_dims) if gather_indices is None: # Drop as many unit-sized dimensions from the transformed shape as we can. gmem_shape = tuple(gmem_ref_ty.shape) for t in gmem_transform: gmem_shape = t.transform_gmem_shape(gmem_shape) # The slice shape may pad along 1-sized dimensions. In that case, we do # not drop them. unit_dims = tuple( i for i, (gs, ss) in enumerate(zip(gmem_shape, slice_shape, strict=True)) if gs == 1 and ss == 1 ) # When issuing an `async_prefetch`, there is no SMEM reference to # transform. if smem_ref is not None: # We need to apply this transform to the SMEM ref as well. The SMEM ref # is expected to have the same shape as the transformed slice shape, # without squeezed dimensions. assert ir.MemRefType(smem_ref.type).shape == slice_shape[num_squeezed_dims:] smem_unit_dims = tuple( i - num_squeezed_dims for i in unit_dims if i - num_squeezed_dims >= 0 ) smem_ref = DropUnitDimsTransform(smem_unit_dims).apply(smem_ref) drop = DropUnitDimsTransform(unit_dims) gmem_transform = (*gmem_transform, drop) dyn_base_indices = drop.transform_index(dyn_base_indices) slice_shape = list(drop.transform_shape(slice_shape)) # After _prepare_async_copy, squeezed dims have been permuted to the # front (via a TransposeTransform in gmem_transform). So in the # transformed shape, they occupy the first `num_squeezed_dims` # positions. # TODO(bchetioui): move the creation of the `TransposeTransform` # here instead of in _prepare_async_copy. squeezed_dims = tuple(d for i, d in enumerate(squeezed_dims) if i not in unit_dims) num_squeezed_dims = len(squeezed_dims) if len(slice_shape) > 5 and squeezed_dims: # We can try to collapse all squeezed dims into one. squeezed_dim_strides = tuple(gmem_strides[d] for d in squeezed_dims) collapse = CollapseLeadingIndicesTransform(squeezed_dim_strides) gmem_transform = (*gmem_transform, collapse) dyn_base_indices = collapse.transform_index(dyn_base_indices) slice_shape = list(collapse.transform_shape(tuple(slice_shape))) num_squeezed_dims = 1 if len(slice_shape) > 5: raise ValueError("Async copies only support striding up to 5 dimensions") del squeezed_dims # pyrefly: ignore[redefinition] dyn_base_indices: list[ir.Value] = list(dyn_base_indices) slice_shape = list(slice_shape) assert all(d == 1 for d in slice_shape[:num_squeezed_dims]) # Partitioned loads have already been processed (before transforms). # We process non-partitioned collective loads here, because only here are we # able to know in what order the data will be written to SMEM. Transposes # and tiling change that order and if we picked a partition based on the # untransformed slice shape, we might have ended up with a non-contiguous # SMEM window, which would no longer be realizable in a single TMA transfer. collective_size = math.prod(self.cluster_size[d] for d in collective) if collective_size > 1 and not isinstance(leader_tracked, _Partitioned): assert gather_indices is None # Checked above. def partition_dim(dim: int, idx: ir.Value, num_chunks: int): # No need to partition squeezed dims. They don't even exist in smem_ref. assert dim >= num_squeezed_dims nonlocal smem_ref slice_shape[dim] //= num_chunks block_offset = arith.muli(idx, c(slice_shape[dim], index)) dyn_base_indices[dim] = arith.addi(dyn_base_indices[dim], block_offset) if smem_ref is not None: smem_ref = utils.memref_slice( smem_ref, (slice(None),) * (dim - num_squeezed_dims) + (utils.ds(block_offset, slice_shape[dim]),), ) idx = utils.cluster_idx(collective) rem_collective_size = collective_size has_swizzle = ( swizzle is not None and swizzle != mgpu_dialect.SwizzlingMode.kNoSwizzle ) # We can partition the minormost dim if there's no swizzling. for dim, slice_size in enumerate( slice_shape[:-1] if has_swizzle else slice_shape ): if slice_size % rem_collective_size == 0: partition_dim(dim, idx, rem_collective_size) rem_collective_size = 1 break elif rem_collective_size % slice_size == 0: # This is an optimization and it lets us skip squeezed dims. if slice_size > 1: dim_idx = arith.remui(idx, c(slice_size, index)) partition_dim(dim, dim_idx, slice_size) idx = arith.divui(idx, c(slice_size, index)) rem_collective_size //= slice_size else: break # We failed to partition the leading dimensions. del idx # We overwrote the block index in the loop. if rem_collective_size > 1: raise ValueError( "None of the leading dimensions in the transformed slice shape" f" {slice_shape} is divisible by the collective size" f" {collective_size}" ) self._check_tma_strides_and_swizzle(gmem_ref_ty, slice_shape, swizzle) return (smem_ref, slice_shape, dyn_base_indices, gmem_transform) def _check_tma_strides_and_swizzle( self, gmem_ref_ty: ir.MemRefType, slice_shape: Sequence[int], swizzle: int | None, ): element_bitwidth = utils.bitwidth(gmem_ref_ty.element_type) gmem_strides, _ = gmem_ref_ty.get_strides_and_offset() if gmem_strides != utils.get_contiguous_strides(gmem_ref_ty.shape): raise NotImplementedError( "Async copies assume the GMEM reference is contiguous" ) if any(s * element_bitwidth % 128 != 0 for s in gmem_strides[:-1]): raise ValueError( "Async copies require all GMEM strides except the last one to be a" " multiple of 16 bytes" ) minor_dim = slice_shape[-1] if slice_shape else 1 if (zeroth_bw := minor_dim * element_bitwidth) % 128 != 0: raise ValueError( "Async copies require the number of bits copied along the last" f" dimension to be divisible by 128, but got {zeroth_bw}" ) if swizzle is not None: swizzle = mgpu_dialect.SwizzlingMode(swizzle) if ( swizzle != mgpu_dialect.SwizzlingMode.kNoSwizzle and minor_dim != (swizzle * 8) // element_bitwidth ): raise ValueError( f"Async copies with {swizzle=} require the last dimension of the" f" slice to be exactly {swizzle} bytes i.e. " f" {(swizzle * 8) // element_bitwidth} elements, but got" f" {minor_dim} elements." ) def async_copy( self, *, src_ref: ir.Value, dst_ref: ir.Value, gmem_slice: Any = (), gmem_transform: MemRefTransform | tuple[MemRefTransform, ...] = (), gmem_peer_id: int | ir.Value | GlobalBroadcast | None = None, barrier: utils.BarrierRef | None = None, swizzle: int | None = None, arrive: bool | None = None, collective: tuple[gpu.Dimension, ...] | gpu.Dimension | None = None, leader_tracked: CopyPartition | None = None, # Should select 0 or 1 threads from the WG. predicate: ir.Value | None | _DefaultPredicate = _DefaultPredicate(), reduction_op: TMAReductionOp | None = None, implementation: AsyncCopyImplementation = AsyncCopyImplementation.TMA, oob_mode: OOBFillMode = OOBFillMode.ZEROS, ): """Initiates an async copy between GMEM and SMEM. Exactly one of ``src_ref`` and ``dst_ref`` must be in GMEM and in SMEM, and the SMEM reference must be contiguous. The GMEM window that is read or written to is specified by ``gmem_slice``. The copy can change the order in which the data appears in the window by applying a sequence of transforms to the GMEM reference (as specified by ``gmem_transform``). When ``collective`` is specified (only allowed for GMEM -> SMEM copies), the identical async_copy must be scheduled by all blocks that share the same coordinates along collective dimensions within a cluster. The behavior is undefined otherwise. The semantics of collective loads depend further on the ``leader_tracked`` argument: - If ``leader_tracked`` is not specified, all blocks load the same data into their shared memory and all receive the update in their barriers, unless ``arrive`` is False. If ``arrive`` is False, you should expect the barrier to have expect_tx incremented by the same amount of bytes as if ``collective`` was not specified. - If ``leader_tracked`` is ``CopyPartition.PARTITIONED(axis)``, each block only loads a separate slice of the data into SMEM, partitioned into equal tiles along the given axis. Only the barrier of the first block in the collective will have its expect_tx incremented by the total size of the transfer across all blocks involved in the collective. Barriers supplied by other blocks will be ignored (even if ``arrive`` is True). - If ``leader_tracked`` is ``CopyPartition.REPLICATED``, all blocks load the same data into their SMEM but only the first block in the collective tracks progress via barrier arrivals. This uses the ``cta_group::2`` mode. """ index = ir.IndexType.get() i8 = ir.IntegerType.get_signless(8) i16 = ir.IntegerType.get_signless(16) i32 = ir.IntegerType.get_signless(32) i64 = ir.IntegerType.get_signless(64) src_ref_ty = ir.MemRefType(src_ref.type) dst_ref_ty = ir.MemRefType(dst_ref.type) element_type = src_ref_ty.element_type element_bitwidth = utils.bitwidth(element_type) if element_type != dst_ref_ty.element_type: raise ValueError( f"Expected same element type, got {element_type} and" f" {dst_ref_ty.element_type}" ) if isinstance(collective, gpu.Dimension): collective = (collective,) elif collective is None: collective = () if not isinstance(gmem_transform, tuple): gmem_transform = (gmem_transform,) if not isinstance(gmem_slice, tuple): gmem_slice = (gmem_slice,) if reduction_op is not None: if implementation != AsyncCopyImplementation.TMA: raise ValueError("Only the TMA implementation supports reductions") if not _is_tma_reduction_op_supported(reduction_op, element_type): raise ValueError( f"Reduction op {reduction_op} not supported by the TMA" f" implementation for element type {element_type}" ) if src_ref_ty.memory_space is None and utils.is_smem_ref(dst_ref_ty): gmem_ref, smem_ref = src_ref, dst_ref if implementation == AsyncCopyImplementation.TMA and barrier is None: raise ValueError("Barriers are required for TMA GMEM -> SMEM copies") if arrive is None: arrive = True # Arrive by default elif utils.is_smem_ref(src_ref_ty) and dst_ref_ty.memory_space is None: gmem_ref, smem_ref = dst_ref, src_ref if barrier is not None: raise ValueError("Barriers are unsupported for SMEM -> GMEM copies") if arrive is None: arrive = True # Commit this copy to the async group by default else: raise ValueError("Only SMEM <-> GMEM copies supported") if collective and gmem_ref is dst_ref: raise ValueError("Only GMEM -> SMEM copies can be collective") ( slice_shape, untransformed_slice_shape, dyn_base_indices, squeezed_dims, gather_indices, gmem_transform, ) = self._prepare_async_copy( gmem_ref, gmem_slice, gmem_transform, collective, leader_tracked, implementation, ) del gmem_slice # Use slice_shape, dyn_base_indices and squeezed_dims instead. gmem_ref_ty = ir.MemRefType(gmem_ref.type) smem_ref_ty = ir.MemRefType(smem_ref.type) # We moved all squeezed dims to the front in _prepare_async_copy. assert all(d == 1 for d in slice_shape[:len(squeezed_dims)]) if slice_shape[len(squeezed_dims):] != smem_ref_ty.shape: raise ValueError( "Expected the SMEM reference to have the same shape as the" f" transformed slice: {tuple(smem_ref_ty.shape)} !=" f" {tuple(slice_shape[len(squeezed_dims):])}" ) if implementation == AsyncCopyImplementation.CP_ASYNC: assert not collective assert leader_tracked is None if not isinstance(predicate, _DefaultPredicate): raise NotImplementedError( "CP_ASYNC needs to be performed by the whole warpgroup and does not" " support the predicate argument" ) # TODO(apaszke): This should be quite easy? The only complication is that # the indices array needs to have a layout compatible with the way we # assign lanes to rows/cols. if gather_indices is not None: raise NotImplementedError("Gather/scatter unsupported for the CP_ASYNC implementation") if smem_ref is src_ref: raise ValueError("CP_ASYNC implementation only supports GMEM -> SMEM copies") offset_scale = 1 if element_bitwidth >= 8 else 8 // element_bitwidth if element_bitwidth < 8: gep_type = i8 elif ( isinstance(element_type, ir.FloatType) and ir.FloatType(element_type).width == 8 ): gep_type = i8 # LLVM has no support for f8. else: gep_type = element_type gmem_strides, _ = gmem_ref_ty.get_strides_and_offset() transformed_strides = gmem_strides for t in gmem_transform: transformed_strides = t.transform_strides(transformed_strides) gmem_offset = utils.dyn_dot( dyn_base_indices, [c(s, index) for s in transformed_strides] ) if offset_scale > 1: gmem_offset = arith.divui(gmem_offset, c(offset_scale, index)) gmem_offset = arith.index_castui(i64, gmem_offset) if squeezed_dims: sliced_dims = [ i for i in range(gmem_ref_ty.rank) if i not in squeezed_dims ] if ( not gmem_transform and sliced_dims and max(squeezed_dims) > min(sliced_dims) ): raise NotImplementedError( "Untiled CP_ASYNC copy expects the GMEM slice to be contiguous in" " memory, so it only supports squeezing leading dimensions" ) # Slice ``gmem_ref`` to drop squeezed dims. gmem_ref = utils.memref_slice( gmem_ref, tuple( 0 if i in squeezed_dims else slice(None) for i in range(gmem_ref_ty.rank) ), ) gmem_ref_ty = gmem_ref.type gmem_strides = [gmem_strides[i] for i in sliced_dims] if not gmem_transform: if ( swizzle is not None and swizzle != mgpu_dialect.SwizzlingMode.kNoSwizzle ): raise NotImplementedError( "Swizzle is not supported for untiled CP_ASYNC copies" ) layout = fa.WGStridedFragLayout.from_shaped_type(smem_ref.type) if layout is None: raise ValueError( "CP_ASYNC requires the number of elements to be a multiple of" f" {fa.WARPGROUP_SIZE}" ) if layout.vec_size * element_bitwidth < 32: raise ValueError( "CP_ASYNC requires at least 4 bytes per transfer" ) gmem_base_ptr = utils.memref_ptr(gmem_ref) gmem_base_ptr = llvm.addrspacecast( llvm.PointerType.get(address_space=1), gmem_base_ptr ) gmem_base_ptr = utils.getelementptr( gmem_base_ptr, [gmem_offset], gep_type ) smem_base_ptr = utils.memref_ptr(smem_ref) bytes_per_transfer = layout.vec_size * element_bitwidth // 8 cache_modifier = ( nvvm.LoadCacheModifierKind.CG if bytes_per_transfer == 16 else nvvm.LoadCacheModifierKind.CA ) for linear_idx in layout.linear_thread_idxs(): idx = arith.index_castui(i32, linear_idx) if offset_scale > 1: idx = arith.divui(idx, c(offset_scale, i32)) smem_ptr = utils.getelementptr(smem_base_ptr, [idx], gep_type) gmem_ptr = utils.getelementptr(gmem_base_ptr, [idx], gep_type) nvvm.cp_async_shared_global( smem_ptr, gmem_ptr, bytes_per_transfer, cache_modifier ) else: assert swizzle is not None swizzle_elems = 8 * swizzle // element_bitwidth tiling = (8, swizzle_elems) if gmem_transform != (TileTransform(tiling),): raise NotImplementedError(gmem_transform) layout = fa.tiled_copy_smem_gmem_layout( *smem_ref_ty.shape[-4:-2], swizzle, element_bitwidth # pyrefly: ignore[bad-argument-count] ) dst_tiled_strides = [ arith.constant(i32, s) for s in layout.tiling.tile_strides(tuple(gmem_strides))[gmem_ref_ty.rank :] ] lane_offset = utils.dyn_dot(layout.lane_indices(), dst_tiled_strides) warp_offset = utils.dyn_dot(layout.warp_indices(), dst_tiled_strides) dyn_offset = arith.addi(lane_offset, warp_offset) dyn_offset = arith.divui(dyn_offset, c(offset_scale, i32)) dyn_offset = arith.extui(i64, dyn_offset) dyn_offset = arith.addi(dyn_offset, gmem_offset) if gmem_ref_ty.rank != 2: raise NotImplementedError("Only 2D copies implemented") gmem_slice_shape = tuple( s for i, s in enumerate(untransformed_slice_shape) if i not in squeezed_dims ) transfers = fa.FragmentedArray.transfer_tiled( smem_ref, swizzle, layout, gmem_slice_shape, optimized=False ) gmem_base_ptr = utils.getelementptr(utils.memref_ptr(gmem_ref), [dyn_offset], gep_type) gmem_base_ptr = llvm.addrspacecast( llvm.PointerType.get(address_space=1), gmem_base_ptr ) bytes_per_transfer = layout.vector_length * element_bitwidth // 8 # Only 16-byte transfers can skip the L1 cache (this is what CG means). cache_modifier = ( nvvm.LoadCacheModifierKind.CG if bytes_per_transfer == 16 else nvvm.LoadCacheModifierKind.CA ) for _get, _update, get_base_idx, smem_ptr in transfers: constant_offset = sum(i * s for i, s in zip(get_base_idx(), gmem_strides, strict=True)) gmem_ptr = utils.getelementptr(gmem_base_ptr, [constant_offset // offset_scale], gep_type) nvvm.cp_async_shared_global(smem_ptr, gmem_ptr, bytes_per_transfer, cache_modifier) if barrier is None: nvvm.cp_async_commit_group() else: # NOTE: Despite its name cp.async.mbarrier.arrive is not an arrival. It # temporarily bumps the pending count (+1 sync, -1 async on completion) # and is overall net-zero. The sole arrival is the ``barrier.arrive`` # called by the leader. The warpgroup barrier ensures that all lanes have # committed their copies before that. nvvm.cp_async_mbarrier_arrive(barrier.get_ptr()) if arrive: utils.warpgroup_barrier() barrier.arrive( predicate=utils.single_thread_predicate(utils.ThreadSubset.WARPGROUP) ) return assert implementation == AsyncCopyImplementation.TMA del smem_ref_ty (smem_ref, slice_shape, dyn_base_indices, gmem_transform) = ( self._prepare_tma( gmem_ref, smem_ref, swizzle, slice_shape, dyn_base_indices, gather_indices, squeezed_dims, gmem_transform, collective, leader_tracked, ) ) assert smem_ref is not None # For type checkers. smem_ref_ty = ir.MemRefType(smem_ref.type) smem_strides, _ = smem_ref_ty.get_strides_and_offset() if any( s != cs and d != 1 # Strides don't matter for dims of size 1. for s, cs, d in zip( smem_strides, utils.get_contiguous_strides(smem_ref_ty.shape), smem_ref_ty.shape, ) ): raise ValueError( "async_copy needs the SMEM reference to be contiguous, but got" f" strides {smem_strides} for shape {smem_ref_ty.shape}" ) collective_size = math.prod(self.cluster_size[d] for d in collective) assert math.prod(slice_shape) * element_bitwidth * collective_size % 8 == 0 transfer_bytes_val = ( math.prod(slice_shape) * element_bitwidth * collective_size // 8 ) # If a copy is multicast, then slice_shape is 1/collective_size of the # local SMEM slice, so each CTA awaits all the bytes that arrive locally. # If leader_tracked is partitioned, then slice_shape corresponds to the # local SMEM slice. We multiply by collective_size to get the total number # of bytes the leader will await. # If leader_tracked is replicated, slice_shape corresponds to 1/collective_size # of the local slice (we use multicast from both blocks). This means that # each CTA will get prod(slice_shape) * collective_size updates, just like # in the multicast case. However, multicast + .cta_group::2 routes all updates # to CTA0, so we need to multiply by collective_size again to get the total. if isinstance(leader_tracked, _Replicated): transfer_bytes_val *= collective_size transfer_bytes = c(transfer_bytes_val, i32) if gather_indices is not None: # Note: Scatters and gathers are handled symmetrically # here. Every mention of "gather" in variable names and the # reasoning below can also be replaced by "scatter" when SMEM is # the source. import builtins zips = functools.partial(builtins.zip, strict=True) # The gather TMA instruction is limited to 2D GMEM references. That means # that we can't apply the transforms to the GMEM reference and have the # TMA engine deal with permuting the data, like we do for non-gather TMA. # Instead, we have to break up the transfer into multiple 2D gathers # ourselves, which requires us to do more complicated stride math etc. # # The minor transformed dim should be a contiguous transfer dim. # The second minor should be a gather dim of size divisible by 4. # The rest can be anything, and we will unroll the transfers over them. if squeezed_dims: raise NotImplementedError("Gather/scatter unsupported when using integer indexing") if reduction_op is not None: raise ValueError("Gather/scatter TMA can't perform reductions") if not isinstance(predicate, _DefaultPredicate): raise ValueError("Gather/scatter TMA can't use a predicate") if gather_indices.layout not in (fa.TMA_INDICES_LAYOUT, fa.TMA_INDICES_4_LAYOUT): raise ValueError(f"Unsupported gather indices layout: {gather_indices.layout}") ROWS_PER_INSTR = 4 # Make sure we'll always be accessing SMEM with sufficient alignment. single_tma_bits = ROWS_PER_INSTR * slice_shape[-1] * element_bitwidth if single_tma_bits % 1024: raise ValueError( "Gather/scatter TMA would require breaking it up into transfers of" f" {single_tma_bits // 8} bytes, but need a multiple of 128 bytes" ) if smem_ref is not src_ref and arrive: assert barrier is not None arrive_predicate = utils.single_thread_predicate(utils.ThreadSubset.WARPGROUP) nvvm.mbarrier_arrive_expect_tx( barrier.get_ptr(), transfer_bytes, predicate=arrive_predicate, ) gmem_strides, _ = gmem_ref_ty.get_strides_and_offset() assert len(gmem_strides) == 2 _, gmem_cols = gmem_ref_ty.shape slice_gather_strides: tuple[int, ...] = (1, 0) # Each row gets a new index, column has no effect. for t in gmem_transform: gmem_strides = t.transform_strides(gmem_strides) slice_gather_strides = t.transform_strides(slice_gather_strides) is_gather_dim = [bool(s) for s in slice_gather_strides] if slice_shape[-1] > 256: raise ValueError( "Gather/scatter TMA can't handle slices with more than 256 columns." " Consider adding a TilingTransform with the minormost dimension <=" " 256." ) tma_desc = self._get_tma_desc( gmem_ref, (), gmem_peer_id, (1, slice_shape[-1]), swizzle, reduction_op, ) assert gather_indices.layout.vector_length == ROWS_PER_INSTR # TMA instructions are uniform, so we can't use multiple lanes. if gather_indices.layout == fa.TMA_INDICES_4_LAYOUT: gather_linear_idx_warp = arith.constant(i32, 0) predicate = utils.single_thread_predicate(utils.ThreadSubset.WARPGROUP) indexing_stride = 1 # Note: If we have several tiles here, we could still issue copies from # each warp to increase parallelism, though it's primarily useful for # saving SMEM when loading a small number of long rows. else: # Indices are split over 4 warps, and replicated within each warp. # Index 00 01 02 03 04 05 06 07 08 09 10 11 12 13 14 15 16 17 ... # Warp <--- 0 ---> <--- 1 ---> <--- 2 ---> <--- 3 ---> <--- 0 -- warp_idx = arith.remui( utils.warp_idx(sync=True), arith.constant(i32, utils.WARPS_IN_WARPGROUP), ) gather_linear_idx_warp = arith.muli(warp_idx, c(ROWS_PER_INSTR, i32)) predicate = utils.single_thread_predicate(utils.ThreadSubset.WARP) indexing_stride = utils.WARPS_IN_WARPGROUP # Since the TMA instruction is limited to 2D gathers, we flatten all # non-gather dims into the column index. max_non_gather_linear_index = sum( (d - 1) * s for g, d, s in zip(is_gather_dim[:-1], slice_shape[:-1], gmem_strides[:-1]) if not g ) # If we ever exceed this then we need to change the size of the GMEM ref, # to prevent the TMA engine from clipping our indices. if max_non_gather_linear_index > gmem_cols: raise NotImplementedError("Non-gather dims don't fit into the columns") col_base_offset = functools.reduce( arith.addi, ( arith.muli(idx, arith.constant(index, stride)) for g, idx, stride in zips( is_gather_dim, dyn_base_indices, gmem_strides ) if not g ), arith.constant(index, 0), ) col_base_offset = arith.index_cast(i32, col_base_offset) # We need to unroll over all non-gather dimensions other than the last one non_gather_slice_shape = tuple( 1 if g else d for d, g in zips(slice_shape[:-1], is_gather_dim[:-1]) ) # First, iterate over gather index registers we have available. for i, reg in enumerate(gather_indices.registers.flat): if utils.bitwidth(gather_indices.mlir_dtype) != 32: reg = arith.extui(ir.VectorType.get((4,), i32), reg) # Compute which rows within the 2D slice we'll be gathering. gather_linear_idx_reg = i * ROWS_PER_INSTR * indexing_stride gather_linear_idx = arith.addi( gather_linear_idx_warp, arith.constant(i32, gather_linear_idx_reg) ) # Transform row indices to align with the transformed SMEM shape. gather_slice_idx = [ arith.remui(arith.divui(gather_linear_idx, c(s, i32)), c(d, i32)) for g, d, s in zip(is_gather_dim, slice_shape, slice_gather_strides) if g ] gather_slice_idx = [arith.index_cast(index, i) for i in gather_slice_idx] gather_rows = [ llvm.extractelement(reg, c(i, i32)) for i in range(ROWS_PER_INSTR) ] # Second, step over non-gather slice indices. for non_gather_idxs in np.ndindex(non_gather_slice_shape): gather_slice_idx_it = iter(gather_slice_idx) smem_indices = tuple( next(gather_slice_idx_it) if g else i for g, i in zip(is_gather_dim[:-1], non_gather_idxs) ) # We should really take a slice here, but it doesn't matter. We're # just going to take the base pointer anyway. transfer_smem_ref = utils.memref_slice(smem_ref, smem_indices) smem_ptr = utils.memref_ptr(transfer_smem_ref) # The slice index needs to be folded into the gather col index. col_slice_offset = sum( idx * stride for g, idx, stride in zips( is_gather_dim[:-1], non_gather_idxs, gmem_strides[:-1] ) if not g ) col_offset = arith.addi(col_base_offset, arith.constant(i32, col_slice_offset)) if smem_ref is src_ref: llvm.inline_asm( ir.Type.parse("!llvm.void"), [predicate, tma_desc, smem_ptr, col_offset, *gather_rows], "@$0 cp.async.bulk.tensor.2d.global.shared::cta.tile::scatter4.bulk_group [$1, {$3, $4, $5, $6, $7}], [$2];", "b,l,r" + ",r" * (ROWS_PER_INSTR + 1), has_side_effects=True, ) else: assert barrier is not None llvm.inline_asm( ir.Type.parse("!llvm.void"), [predicate, smem_ptr, tma_desc, barrier.get_ptr(), col_offset, *gather_rows], "@$0 cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes [$1], [$2, {$4, $5, $6, $7, $8}], [$3];", "b,r,l,r" + ",r" * (ROWS_PER_INSTR + 1), has_side_effects=True, ) if smem_ref is src_ref and arrive: nvvm.cp_async_bulk_commit_group() return assert gather_indices is None # Only tiled TMA handled below. smem_ptr = utils.memref_ptr(smem_ref) if isinstance(predicate, _DefaultPredicate): predicate = utils.single_thread_predicate(utils.ThreadSubset.WARPGROUP) if predicate is None: predicate = c(1, ir.IntegerType.get_signless(1)) def check_contiguous_slice(slice_shape, strides): if strides[-1] != 1: return False expected_stride = 1 for dim, stride in zip(reversed(slice_shape), reversed(strides), strict=True): if dim != 1 and stride != expected_stride: return False expected_stride *= dim return True ref = gmem_ref for t in gmem_transform: ref = t.apply(ref) ref_ty = ir.MemRefType(ref.type) strides, _ = ref_ty.get_strides_and_offset() # Use the simpler copy instruction for bulk contiguous transfers. is_simple_contiguous_copy = ( check_contiguous_slice(slice_shape, strides) and reduction_op is None and ( swizzle is None or swizzle == mgpu_dialect.SwizzlingMode.kNoSwizzle ) and collective_size == 1 and oob_mode != OOBFillMode.ZEROS and leader_tracked is None and implementation == AsyncCopyImplementation.TMA ) if is_simple_contiguous_copy: slicing = tuple( utils.ds(start, size) for start, size in zip(dyn_base_indices, slice_shape) ) ref_slice = utils.memref_slice(ref, slicing) del ref # Use ref_slice from now on. gmem_base_ptr = utils.memref_ptr(ref_slice) if oob_mode == OOBFillMode.PROMISE_IN_BOUNDS: clamped_transfer_bytes = transfer_bytes else: index = ir.IndexType.get() to_i32 = lambda x: arith.index_cast(ir.IntegerType.get_signless(32), x) # Clamp the out-of-bounds accesses to be within the bounds of the GMEM ref shape. total_elements = c(math.prod(gmem_ref_ty.shape), index) linear_offset_elems = memref.extract_strided_metadata(ref_slice)[1] # pyrefly: ignore[bad-index] # Equivalent to: max(0, total - offset) valid_offset_elems = arith.minui(linear_offset_elems, total_elements) remaining_elements = arith.subi(total_elements, valid_offset_elems) remaining_bytes = utils.elements_to_bytes(remaining_elements, element_bitwidth) transfer_bytes_idx = arith.index_cast(index, transfer_bytes) clamped_transfer_bytes = to_i32(arith.minui(transfer_bytes_idx, remaining_bytes)) if gmem_peer_id is not None: if isinstance(gmem_peer_id, GlobalBroadcast): multimem_ref = self.to_remote_multicast(ref_slice, on_host=False) gmem_base_ptr = utils.memref_ptr(multimem_ref.ref) else: if isinstance(gmem_peer_id, int): gmem_peer_id = c(gmem_peer_id, i32) remote_ref = self.to_remote(ref_slice, gmem_peer_id, on_host=False) gmem_base_ptr = utils.memref_ptr(remote_ref) if gmem_ref is src_ref: assert barrier is not None barrier_ptr = barrier.get_ptr() if arrive: nvvm.mbarrier_arrive_expect_tx( barrier_ptr, clamped_transfer_bytes, predicate=predicate ) else: if oob_mode != OOBFillMode.PROMISE_IN_BOUNDS: # Make sure the barrier is completed with the correct number of bytes # if we resize the transfer size. left_to_complete = arith.subi(transfer_bytes, clamped_transfer_bytes) barrier.complete_tx(left_to_complete, predicate=predicate) llvm.inline_asm( ir.Type.parse("!llvm.void"), [predicate, smem_ptr, gmem_base_ptr, clamped_transfer_bytes, barrier_ptr], """ @$0 cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [$1], [$2], $3, [$4]; """, "b,l,l,r,l", has_side_effects=True, ) else: llvm.inline_asm( ir.Type.parse("!llvm.void"), [predicate, gmem_base_ptr, smem_ptr, clamped_transfer_bytes], """ @$0 cp.async.bulk.global.shared::cta.bulk_group [$1], [$2], $3; """, "b,l,l,r", has_side_effects=True, ) if arrive: nvvm.cp_async_bulk_commit_group() return # Below are tiled TMA copies using a tensormap. if max(slice_shape) > 256: raise ValueError( "Async copies only support copying <=256 elements along each" f" dimension, got {tuple(slice_shape)}" ) tma_desc = self._get_tma_desc( gmem_ref, gmem_transform, gmem_peer_id, tuple(slice_shape), swizzle, reduction_op, ) # We construct TMA descriptors in column-major order. rev_dyn_base_indices = [ arith.index_cast(i32, idx) for idx in reversed(dyn_base_indices) ] if gmem_ref is src_ref: assert barrier is not None barrier_ptr = barrier.get_ptr() assert reduction_op is None rank = len(slice_shape) idx_operands = ",".join(f"${i}" for i in range(4, 4 + rank)) if collective_size > 1 and leader_tracked is not None: assert collective_size == 2 if arrive: first_block = arith.cmpi( arith.CmpIPredicate.eq, utils.cluster_idx(collective), c(0, index), ) arrive_predicate = arith.andi(predicate, first_block) nvvm.mbarrier_arrive_expect_tx( barrier_ptr, transfer_bytes, predicate=arrive_predicate ) if isinstance(leader_tracked, _Replicated): multicast_mask = (arith.trunci( i16, utils.cluster_collective_mask(self.cluster_size, collective) ),) smem_space = "shared::cluster" multicast_mod = ".multicast::cluster" multicast_operand = f", ${4 + rank}" else: multicast_mask = () smem_space = "shared::cta" multicast_mod = "" multicast_operand = "" mapped_barrier_ptr = barrier.remap_to_cluster( gpu.Dimension.x, c(0, index) ).get_ptr() llvm.inline_asm( ir.Type.parse("!llvm.void"), [predicate, smem_ptr, tma_desc, mapped_barrier_ptr, *rev_dyn_base_indices, *multicast_mask], f""" @$0 cp.async.bulk.tensor.{rank}d.{smem_space}.global.tile.mbarrier::complete_tx::bytes{multicast_mod}.cta_group::2 [$1], [$2, {{{idx_operands}}}], [$3]{multicast_operand}; """, "b,r,l,r" + ",r" * rank + ",h" * len(multicast_mask), has_side_effects=True, ) else: if arrive: nvvm.mbarrier_arrive_expect_tx( barrier_ptr, transfer_bytes, predicate=predicate ) if collective_size > 1: multicast_mask = (arith.trunci( i16, utils.cluster_collective_mask(self.cluster_size, collective) ),) smem_space = "shared::cluster" multicast_mod = ".multicast::cluster" multicast_operand = f", ${4 + rank}" else: multicast_mask = () smem_space = "shared::cta" multicast_mod = "" multicast_operand = "" llvm.inline_asm( ir.Type.parse("!llvm.void"), [predicate, smem_ptr, tma_desc, barrier_ptr, *rev_dyn_base_indices, *multicast_mask], ( f"@$0 cp.async.bulk.tensor.{rank}d" f".{smem_space}.global.mbarrier::complete_tx::bytes{multicast_mod} " f"[$1], [$2, {{{idx_operands}}}], [$3]{multicast_operand};" ), "b,r,l,r" + ",r" * rank + ",h" * len(multicast_mask), has_side_effects=True, ) else: if reduction_op is not None: rank = len(slice_shape) idx_operands = ",".join(f"${i}" for i in range(3, 3 + rank)) llvm.inline_asm( ir.Type.parse("!llvm.void"), [predicate,smem_ptr,tma_desc,*rev_dyn_base_indices], f"@$0 cp.reduce.async.bulk.tensor.{rank}d.global.shared::cta.{_reduction_op_to_ptx(reduction_op)}.tile.bulk_group [$2,{{{idx_operands}}}], [$1];", "b,r,l" + ",r" * rank, has_side_effects=True, ) if arrive: nvvm.cp_async_bulk_commit_group() else: nvvm.cp_async_bulk_tensor_global_shared_cta( tma_desc, smem_ptr, rev_dyn_base_indices, predicate=predicate ) if arrive: nvvm.cp_async_bulk_commit_group() def async_copy_im2col( self, *, src_ref: ir.Value, dst_ref: ir.Value, barrier: utils.BarrierRef, window_start_bounds: Sequence[tuple[int, int]], start_indices: Sequence[int | ir.Value], filter_offsets: Sequence[int | ir.Value], window_strides: Sequence[int] | None = None, swizzle: int | None = None, arrive: bool = True, # Should select 0 or 1 threads from the WG. predicate: ir.Value | None | _DefaultPredicate = _DefaultPredicate(), ): """Initiates an async im2col copy from GMEM to SMEM. This is the building block of implicit-GEMM convolutions. Each row of ``dst_ref`` receives the channels of one pixel of one window in ``src_ref``. - ``src_ref`` is a batch of channels-last images of shape ``(N, W, C)``, ``(N, H, W, C)`` or ``(N, D, H, W, C)``. - ``dst_ref`` has shape ``(*rows, cols)``. Its leading dimensions are flattened into ``math.prod(rows)`` windows of ``cols`` channels each. - ``window_start_bounds`` gives, for each spatial dimension of size ``S``, a ``(lo, hi)`` pair: windows start at positions ``lo``, ``lo + s``, ``lo + 2 * s``, ... below ``S + hi``, where ``s`` is the corresponding window stride, with the image implicitly zero-padded outside ``[0, S)``. For example, a 3x3 ``SAME`` convolution uses ``(-1, -1)`` in every spatial dimension. - ``start_indices`` is ``(n, *start, c)``: the batch index and start position of the window loaded into the first row, and the first channel to load. Each subsequent row takes the next window in row-major order over ``(n, *start)``, so a copy can span several images. Rows past the last image are zero-filled. - ``filter_offsets`` selects the pixel to load within each window, as an offset from the window's start position. Offsets must be in ``[0, 2**(16 // (rank - 2)) - 1]``, i.e. ``[0, 65535]`` for rank 3, ``[0, 255]`` for rank 4 and ``[0, 31]`` for rank 5. Only static offsets are checked; out-of-range dynamic offsets give unspecified results. - ``window_strides`` gives, for each spatial dimension, the step between consecutive window start positions. Each stride must be in ``[1, 8]``; defaults to 1 in every spatial dimension. ``barrier``, ``swizzle``, ``arrive`` and ``predicate`` work as in ``async_copy``. """ # TODO(slebedev): Remove once the minimum jaxlib version is 0.12.0. if jaxlib_version < (0, 12, 0): raise NotImplementedError("async_copy_im2col requires jaxlib >=0.12.0") c = utils.c index = ir.IndexType.get() i16 = ir.IntegerType.get_signless(16) i32 = ir.IntegerType.get_signless(32) gmem_ref, smem_ref = src_ref, dst_ref gmem_ref_ty = ir.MemRefType(gmem_ref.type) smem_ref_ty = ir.MemRefType(smem_ref.type) if gmem_ref_ty.memory_space is not None or not utils.is_smem_ref(smem_ref_ty): raise ValueError( "async_copy_im2col requires src_ref to be in GMEM and dst_ref to be" " in SMEM" ) if gmem_ref_ty.element_type != smem_ref_ty.element_type: raise ValueError( f"Expected same element type, got {gmem_ref_ty.element_type} and" f" {smem_ref_ty.element_type}" ) smem_strides, _ = smem_ref_ty.get_strides_and_offset() if any( s != cs and d != 1 # Strides don't matter for dims of size 1. for s, cs, d in zip( smem_strides, utils.get_contiguous_strides(smem_ref_ty.shape), smem_ref_ty.shape, strict=True, ) ): raise ValueError( "async_copy_im2col needs the SMEM reference to be contiguous, but got" f" strides {smem_strides} for shape {smem_ref_ty.shape}" ) element_bw = utils.bitwidth(gmem_ref_ty.element_type) if element_bw < 8: raise ValueError( "im2col TMA does not support sub-byte types, got" f" {element_bw}-bit elements" ) window_start_bounds = tuple( (int(lo), int(hi)) for lo, hi in window_start_bounds ) rank = gmem_ref_ty.rank if rank not in (3, 4, 5): raise ValueError(f"im2col TMA requires GMEM rank in (3, 4, 5), got {rank}") num_spatial = rank - 2 if len(window_start_bounds) != num_spatial: raise ValueError( f"Expected {num_spatial} (lo, hi) pairs in window_start_bounds, got" f" {len(window_start_bounds)}" ) bits = 16 // num_spatial allowed_bounds = range(-(1 << (bits - 1)), 1 << (bits - 1)) for i, (lower, upper) in enumerate(window_start_bounds): if lower not in allowed_bounds or upper not in allowed_bounds: raise ValueError( f"window_start_bounds along spatial dim {i} must be in" f" [{allowed_bounds.start}, {allowed_bounds.stop - 1}] for rank" f" {rank}, but got ({lower}, {upper})" ) if gmem_ref_ty.shape[i + 1] + upper - lower < 1: raise ValueError( f"window_start_bounds along spatial dim {i} give an empty range" f" [{lower}, {gmem_ref_ty.shape[i + 1] + upper}) of window start" " positions" ) if window_strides is None: window_strides = (1,) * num_spatial window_strides = tuple(window_strides) if len(window_strides) != num_spatial: raise ValueError( f"Expected {num_spatial} window_strides, got {len(window_strides)}" ) for i, stride in enumerate(window_strides): if not (1 <= stride <= 8): raise ValueError( f"window_strides along spatial dim {i} must be in [1, 8], but got" f" {stride}" ) if smem_ref_ty.rank < 2: raise ValueError( "async_copy_im2col requires dst_ref to have rank >= 2, got" f" {smem_ref_ty.rank}" ) num_rows = math.prod(smem_ref_ty.shape[:-1]) num_channels = smem_ref_ty.shape[-1] if not (1 <= num_rows <= 1024): raise ValueError( f"dst_ref must have between 1 and 1024 rows, got {num_rows}" ) if not (1 <= num_channels <= 256): raise ValueError( f"dst_ref must have between 1 and 256 columns, got {num_channels}" ) self._check_tma_strides_and_swizzle( gmem_ref_ty, (num_rows, num_channels), swizzle ) if len(start_indices) != rank: raise ValueError( f"Expected {rank} start_indices, got {len(start_indices)}" ) dyn_start_indices: list[ir.Value] = [] for idx in start_indices: if isinstance(idx, (ir.Operation, ir.OpView)): idx = idx.result if isinstance(idx, int): dyn_start_indices.append(c(idx, i32)) elif isinstance(idx, ir.Value): if idx.type != index: raise ValueError( f"Dynamic start_indices must have index type, got {idx.type}" ) dyn_start_indices.append(arith.index_cast(i32, idx)) else: raise TypeError(f"Unsupported start index type: {type(idx)}") if len(filter_offsets) != num_spatial: raise ValueError( f"Expected {num_spatial} filter_offsets, got {len(filter_offsets)}" ) allowed_offsets = range(1 << bits) dyn_filter_offsets: list[ir.Value] = [] for offset in filter_offsets: if isinstance(offset, (ir.Operation, ir.OpView)): offset = offset.result if isinstance(offset, int): if offset not in allowed_offsets: raise ValueError( f"filter_offsets must be in [0, {allowed_offsets.stop - 1}]" f" for rank {rank}, got {offset}" ) dyn_filter_offsets.append(c(offset, i16)) elif isinstance(offset, ir.Value): if offset.type != index: raise ValueError( f"Dynamic filter_offsets must have index type, got {offset.type}" ) dyn_filter_offsets.append(arith.index_cast(i16, offset)) else: raise TypeError(f"Unsupported filter offset type: {type(offset)}") assert math.prod(smem_ref_ty.shape) * element_bw % 8 == 0 transfer_bytes = c(math.prod(smem_ref_ty.shape) * element_bw // 8, i32) tma_desc = self._get_tma_desc( gmem_ref, (), None, (num_rows, num_channels), swizzle, reduction_op=None, window_start_bounds=window_start_bounds, window_strides=window_strides, ) barrier_ptr = barrier.get_ptr() if isinstance(predicate, _DefaultPredicate): predicate = utils.single_thread_predicate(utils.ThreadSubset.WARPGROUP) if predicate is None: predicate = c(1, ir.IntegerType.get_signless(1)) if arrive: nvvm.mbarrier_arrive_expect_tx( barrier_ptr, transfer_bytes, predicate=predicate ) idx_operands = ", ".join(f"${i}" for i in range(3, 3 + rank)) offset_operands = ", ".join( f"${i}" for i in range(3 + rank, 3 + rank + num_spatial) ) utils.inline_ptx( f"cp.async.bulk.tensor.{rank}d" ".shared::cta.global.im2col.mbarrier::complete_tx::bytes" f" [$0], [$1, {{{idx_operands}}}], [$2], {{{offset_operands}}};", utils.memref_ptr(smem_ref), tma_desc, barrier_ptr, *reversed(dyn_start_indices), *reversed(dyn_filter_offsets), predicate=predicate, ) def async_prefetch( self, *, gmem_ref: ir.Value, gmem_slice: Any = (), gmem_transform: MemRefTransform | tuple[MemRefTransform, ...] = (), gmem_peer_id: int | ir.Value | None = None, swizzle: int | None = None, collective: tuple[gpu.Dimension, ...] | gpu.Dimension | None = None, leader_tracked: CopyPartition | None = None, # Should select 0 or 1 threads from the WG. predicate: ir.Value | None | _DefaultPredicate = _DefaultPredicate(), ): i32 = ir.IntegerType.get_signless(32) if isinstance(collective, gpu.Dimension): collective = (collective,) elif collective is None: collective = () else: collective = tuple(collective) if not isinstance(gmem_transform, tuple): gmem_transform = (gmem_transform,) if not isinstance(gmem_slice, tuple): gmem_slice = (gmem_slice,) impl = AsyncCopyImplementation.TMA ( slice_shape, _, dyn_base_indices, squeezed_dims, gather_indices, gmem_transform, ) = self._prepare_async_copy( gmem_ref, gmem_slice, gmem_transform, collective, leader_tracked, impl ) del gmem_slice # Use slice_shape, dyn_base_indices and squeezed_dims instead. (_, slice_shape, dyn_base_indices, gmem_transform) = ( self._prepare_tma( gmem_ref, None, swizzle, slice_shape, dyn_base_indices, gather_indices, squeezed_dims, gmem_transform, collective, leader_tracked, ) ) if gather_indices is not None: raise NotImplementedError("Gather/scatter prefetch not implemented yet") tma_desc = self._get_tma_desc( gmem_ref, gmem_transform, gmem_peer_id, tuple(slice_shape), swizzle, reduction_op=None, ) # We construct TMA descriptors in column-major order. rev_dyn_base_indices = [ arith.index_cast(i32, idx) for idx in reversed(dyn_base_indices) ] if isinstance(predicate, _DefaultPredicate): predicate = utils.single_thread_predicate(utils.ThreadSubset.WARPGROUP) if predicate is None: predicate = c(1, ir.IntegerType.get_signless(1)) rank = len(slice_shape) idx_operands = ",".join(f"${i}" for i in range(2, 2 + rank)) llvm.inline_asm( ir.Type.parse("!llvm.void"), [predicate, tma_desc, *rev_dyn_base_indices], f"@$0 cp.async.bulk.prefetch.tensor.{rank}d.L2.global.tile [$1, {{{idx_operands}}}];", "b,l" + ",r" * rank, has_side_effects=True, ) def await_async_copy( self, allow_groups: int, await_read_only: bool = False, scope: utils.ThreadSubset = utils.ThreadSubset.WARPGROUP, ): nvvm.cp_async_bulk_wait_group(allow_groups, read=await_read_only) if scope == utils.ThreadSubset.WARPGROUP: utils.warpgroup_barrier() elif scope == utils.ThreadSubset.WARP: utils.warp_barrier() else: raise ValueError(f"Unsupported scope: {scope}") def await_cp_async_copy(self, allow_groups: int): nvvm.cp_async_wait_group(allow_groups) utils.warpgroup_barrier() def _ensure_nvshmem_decls(self): if self.is_device_collective or self.device_collective_metadata is not None: return self.is_device_collective = True with ir.InsertionPoint(self.module.body): nvshmem_my_pe_type = ir.TypeAttr.get(ir.Type.parse("!llvm.func")) llvm.LLVMFuncOp( "nvshmem_my_pe", nvshmem_my_pe_type, sym_visibility="private" ) nvshmem_ptr_type = ir.TypeAttr.get( ir.Type.parse("!llvm.func") ) llvm.LLVMFuncOp("nvshmem_ptr", nvshmem_ptr_type, sym_visibility="private") nvshmemx_mc_ptr_type = ir.TypeAttr.get( ir.Type.parse("!llvm.func") ) llvm.LLVMFuncOp( "nvshmemx_mc_ptr", nvshmemx_mc_ptr_type, sym_visibility="private" ) def _find_kernel_argument_index(self, ref: ir.Value): """Finds the index of the kernel argument used to derive the given reference.""" if not isinstance(ref.type, ir.MemRefType): raise ValueError(f"Expected a memref, got {ref.type}") op = ref.owner while KERNEL_ARG_ID_ATTR not in getattr(op, "attributes", {}): if not isinstance(op, ir.OpView): raise ValueError( f"Can't find the kernel argument for the reference: {ref}" ) memref_operands = [ operand for operand in op.operands if isinstance(operand.type, ir.MemRefType) ] if len(memref_operands) != 1: raise ValueError( f"Can't find the kernel argument. {op} doesn't have a single memref" " operand." ) op = memref_operands[0].owner attr = op.attributes[KERNEL_ARG_ID_ATTR] # pyrefly: ignore[missing-attribute] # Save the result so that we can find it out faster next time. ref.owner.attributes[KERNEL_ARG_ID_ATTR] = attr # pyrefly: ignore[missing-attribute] return attr.value def _get_collective_metadata(self, on_host: bool): return ( self.host_collective_metadata if on_host else self.device_collective_metadata ) def _get_parameter_address_on_peer( self, parameter_id: int, peer_id: ir.Value, on_host: bool ): """Returns the address of the given parameter on the given peer.""" collective_metadata = self._get_collective_metadata(on_host) index = ir.IndexType.get() offset_to_parameter = arith.constant( ir.IndexType.get(), COLLECTIVE_METADATA_SIZE + self.num_peers * parameter_id, ) assert peer_id.type == ir.IntegerType.get_signless(32) parameter_on_device_offset = arith.addi( offset_to_parameter, arith.index_cast(index, peer_id) ) return memref.load(collective_metadata, [parameter_on_device_offset]) def _get_offset_to_parameter( self, ref: ir.Value, parameter_address: ir.Value ): ref_int = llvm.ptrtoint(ir.IntegerType.get_signless(64), ref) return arith.subi(ref_int, parameter_address) def _mark_parameters_if_multiprocess(self): # All multi-process parameters should be allocated in collective memory. if self.multihost_kernel: symmetric_memory_parameters = np.ones(self.num_params, dtype=np.bool) self.module.operation.attributes[SYMMETRIC_MEMORY_ARGS_ATTR] = ( ir.DenseIntElementsAttr.get(symmetric_memory_parameters) ) def to_remote( self, ref: ir.Value, peer: ir.Value, *, _kernel_arg_idx: int | None = None, on_host: bool = False, ): collective_metadata = self._get_collective_metadata(on_host) i32 = ir.IntegerType.get_signless(32) if isinstance(ref.type, ir.MemRefType): assert _kernel_arg_idx is None ref_ty = ir.MemRefType(ref.type) ptr_ty = llvm.PointerType.get() desc = utils.to_memref_descriptor(ref) arg_idx = None if collective_metadata is not None: arg_idx = self._find_kernel_argument_index(ref) # Extract the aligned base pointer, as we want to annotate it with the # relevant alignment information. aligned_ptr = llvm.extractvalue(ptr_ty, desc, [1]) remote_ptr = self.to_remote( aligned_ptr, peer, _kernel_arg_idx=arg_idx, on_host=on_host, ) desc = llvm.insertvalue(desc, remote_ptr, [0]) desc = llvm.insertvalue(desc, remote_ptr, [1]) remote_memref = builtin.unrealized_conversion_cast([ref_ty], [desc]) assert isinstance(remote_memref, ir.Value) if collective_metadata is not None: remote_memref.owner.attributes[KERNEL_ARG_ID_ATTR] = ir.IntegerAttr.get( i32, arg_idx ) return remote_memref if collective_metadata is None: self._ensure_nvshmem_decls() if ref.type != llvm.PointerType.get(): raise ValueError(f"Unsupported type for to_remote: {ref.type}") if peer.type != i32: raise ValueError(f"peer index must be an i32, got {peer.type}") return llvm.call(ref.type, [ref, peer], [], [], callee="nvshmem_ptr") else: self._mark_parameters_if_multiprocess() # Collective metadata contains pointers of kernel arguments for each peer # device. The pointer has the following format: # [ # param0_peer0, param0_peer1, ..., param0_peerN, # param1_peer0, param1_peer1, ..., param1_peerN, # ... # ] # During the lowering we need to find the corresponding kernel argument # for a given reference, load the corresponding pointer from the # collective metadata and also compute the address of the given reference. # As an example an address of signlas will have an offset from the first # pointer of the kernel arguments defined with the memref.subview # operation. self.module.operation.attributes[COLLECTIVE_ATTR] = ir.UnitAttr.get() assert _kernel_arg_idx is not None # TODO(apaszke): Just use the pointer directly. After all it is an arg. current_device = self.device_id(on_host) parameter_on_current_device = self._get_parameter_address_on_peer( _kernel_arg_idx, current_device, on_host ) parameter_on_peer_device = self._get_parameter_address_on_peer( _kernel_arg_idx, peer, on_host ) peer_offset = arith.subi( parameter_on_peer_device, parameter_on_current_device ) i64 = ir.IntegerType.get_signless(64) # Enforce alignment to allow vectorized loads and stores to be generated # whenever possible. is_aligned = arith.cmpi( arith.CmpIPredicate.eq, arith.andi(peer_offset, c(HBM_ALIGNMENT - 1, i64)), c(0, i64), ) llvm.intr_assume(is_aligned, [], ir.DenseI32ArrayAttr.get([])) return utils.getelementptr( ref, [peer_offset], ir.IntegerType.get_signless(8) ) def to_remote_multicast(self, ref: ir.Value, on_host: bool = False): i32 = ir.IntegerType.get_signless(32) self._flag_multimem_usage() if not isinstance(ref.type, ir.MemRefType): raise ValueError(f"Unsupported type for to_remote_multicast: {ref.type}") # We replace the offset in the ref type by 0, because memref_ptr always # folds the offset into the pointer. ref_ty = ir.MemRefType(ref.type) strides, _ = ref_ty.get_strides_and_offset() result_type = ir.MemRefType.get( ref_ty.shape, ref_ty.element_type, ir.StridedLayoutAttr.get(0, strides), ref_ty.memory_space, ) collective_metadata = self._get_collective_metadata(on_host) if collective_metadata is None: self._ensure_nvshmem_decls() world_team = arith.constant(i32, 0) ptr = utils.memref_ptr(ref) mc_ptr = llvm.call( ptr.type, [world_team, ptr], [], [], callee="nvshmemx_mc_ptr" ) return utils.MultimemRef(utils.ptr_as_memref(mc_ptr, result_type)) parameter_id = self._find_kernel_argument_index(ref) # In multiprocess mode, all parameters should be allocated within collective # memory. module_attributes = self.module.operation.attributes self._mark_parameters_if_multiprocess() if not self.multihost_kernel: if SYMMETRIC_MEMORY_ARGS_ATTR in module_attributes: symmetric_memory_parameters = np.array( module_attributes[SYMMETRIC_MEMORY_ARGS_ATTR] ) else: symmetric_memory_parameters = np.zeros(self.num_params, dtype=np.bool) symmetric_memory_parameters[parameter_id] = True module_attributes[SYMMETRIC_MEMORY_ARGS_ATTR] = ( ir.DenseIntElementsAttr.get(symmetric_memory_parameters) ) current_device = self.device_id(on_host) parameter_on_current_device = self._get_parameter_address_on_peer( parameter_id, current_device, on_host ) # Load the multimem address of the parameter from the collective metadata. multimem_addresses_offset = ( COLLECTIVE_METADATA_SIZE + self.num_peers * self.num_params + parameter_id ) multimem_ptr_offset = arith.constant( ir.IndexType.get(), multimem_addresses_offset ) multimem_ptr = memref.load(collective_metadata, [multimem_ptr_offset]) # Compute the offset of the given reference to the parameter and add it to # the multimem address of the parameter. ref_offset = self._get_offset_to_parameter( utils.memref_ptr(ref), parameter_on_current_device ) multimem_address = arith.addi(multimem_ptr, ref_offset) multimem_ptr = llvm.inttoptr(llvm.PointerType.get(), multimem_address) return utils.MultimemRef(utils.ptr_as_memref(multimem_ptr, result_type)) def device_id(self, on_host: bool = False) -> ir.Value: collective_metadata = self._get_collective_metadata(on_host) i32 = ir.IntegerType.get_signless(32) if collective_metadata is None: self._ensure_nvshmem_decls() return cast(ir.Value, llvm.call(i32, [], [], [], callee="nvshmem_my_pe")) else: self._mark_parameters_if_multiprocess() # Rank id is stored as the first element of the collective metadata. self.module.operation.attributes[COLLECTIVE_ATTR] = ir.UnitAttr.get() rank_offset_constant = arith.constant(ir.IndexType.get(), 0) load_op = memref.load(collective_metadata, [rank_offset_constant]) device_rank = arith.trunci(i32, load_op) device_rank.owner.attributes[DEVICE_ID_ATTR] = ir.UnitAttr.get() return device_rank class ReplicationError(Exception): pass