"""CuTeDSL override registrations for ``aten::scatter_add``. Two kernels for the fast-path layout (self/src 2D-shaped after coalescing, with self contiguous on the slice axis and src either contiguous or row-strided), tried in order: 1. **TMA** (``tma_kernel.py``, sm_90+): uses ``cp.reduce.async.bulk`` to offload the whole reduction to the TMA unit, 3-7x faster than aten on high-contention embedding-bag-grad-style workloads. 2. **vec-scatter** (``vec_scatter_kernel.py``, any SM): warp-per-row scheduling with ``atomicAdd`` (scalar fp32, paired x2 halves). Covers pre-sm_90, and sm_90+ shapes where the TMA alignment constraint isn't met. Eligibility mirrors aten's ``fast_scatter_add_kernel_eligible`` (IndexKernelUtils.h): we restride ``self`` so its scatter-axis stride is 0 and shape matches ``index``, then build a ``TensorIterator`` on ``(self_restrided, src_restrided, index)``. TensorIterator's coalescing and stride-magnitude reordering produces a 2D iterator iff the layout is kernel-friendly; we then pattern-match the resulting iter strides to confirm slice/index-axis assignment. This delegates the layout analysis to ``TensorIterator`` rather than reimplementing it. Anything else (non-2D-coalescing layout, dtype mismatch, deterministic mode) falls through to aten. Aten is competitive or better on those shapes, so we intentionally don't override them. In-place (``scatter_add_``) is registered explicitly; PyTorch's structured-delegate from ``scatter_add_`` to ``scatter_add.out`` happens below the dispatcher, so overriding the ``.out`` variant alone doesn't intercept the in-place method. """ import functools import torch from torch._tensor_iterator import TensorIterator from ... import cutedsl_utils as cu from ...registry import _OpCondFn, _OpImplFn _SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32} # 16B alignment for `src.data_ptr()` and `src.stride(0) * elem`: required # for the LDG.128 / tile-mode TMA load of `src` and the row-start 16B # boundary that LDG.128 reads from. Applies to both paths. _SRC_ALIGN_BYTES = 16 # Path-specific destination alignment. TMA: cp.reduce.async.bulk gmem # operand needs 16B base + stride. vec-scatter: scalar fp32 atomicAdd / # f16x2 / bf16x2 paired atomics need natural alignment of the 4B atomic # operand on both base and per-row stride. _TMA_DST_ALIGN_BYTES = 16 _VEC_DST_ALIGN_BYTES = 4 @functools.cache def _is_rocm() -> bool: # These kernels are built on the CUTLASS Python DSL (CuteDSL), which is # NVIDIA-only and performs a hard CUDA-driver dependency check at compile # time. On a ROCm/HIP build ``tensor.is_cuda`` is still True, so without # this guard the cond approves and ``cute.compile`` fails with # CudaDriverDependencyError. Decline so scatter_add falls back to aten. return torch.version.hip is not None @functools.cache def _has_sm90_plus() -> bool: if not torch.cuda.is_available(): return False major, _ = torch.cuda.get_device_capability(0) return major >= 9 def _deterministic() -> bool: # Our kernels use non-deterministic atomicAdd. Under # torch.use_deterministic_algorithms(True), aten's scatter_add falls # into _scatter_via_index_put instead; we must decline to preserve # that behavior. return torch.are_deterministic_algorithms_enabled() def _any_cow(*tensors: torch.Tensor) -> bool: return any( torch._C._is_cow_tensor(t) # pyrefly: ignore[missing-attribute] for t in tensors ) def _base_cond_ok(*tensors: torch.Tensor) -> bool: """Pre-checks shared by every path: env sanity, all CUDA, non-COW.""" if _is_rocm(): return False if _deterministic(): return False if not all(t.is_cuda for t in tensors): return False if _any_cow(*tensors): return False return True # --------------------------------------------------------------------------- # TI-driven layout analysis. # # Mirrors aten's pattern in ScatterGatherKernel.cu / IndexKernelUtils.h: # restride ``self`` so its scatter-axis stride is 0 and its shape matches # ``index``, then build a TensorIterator on the three. TensorIterator does # the broadcast / coalesce / dim reorder; we then pattern-match the post- # build iter strides to recognize the kernel's expected layout. # --------------------------------------------------------------------------- def _normalize_dim(dim: int, ndim: int) -> int: return dim + ndim if dim < 0 else dim def _scatter_add_eligibility( self: torch.Tensor, d: int, index: torch.Tensor, src: torch.Tensor ) -> TensorIterator | None: """Return the analysis TensorIterator if (self, d, index, src) fits the kernel's expected layout, else ``None``. ``d`` is the already-normalized scatter axis. Mirrors aten's ``fast_scatter_add_kernel_eligible`` (IndexKernelUtils.h): restride ``self`` so its scatter-axis stride is 0 (shape = ``index.shape``), let TensorIterator coalesce + reorder, then check that the result is a 2D iter where dim 0 is the contiguous slice axis and dim 1 is the scatter (index) axis. """ if self.dtype not in _SUPPORTED_DTYPES or self.dtype != src.dtype: return None if index.dtype not in (torch.int32, torch.int64): return None if self.ndim != src.ndim or self.ndim != index.ndim or self.ndim == 0: return None if not 0 <= d < self.ndim: return None self_strides = list(self.stride()) self_strides[d] = 0 # aten's restride_dim trick try: self_r = self.as_strided(index.shape, self_strides) src_r = src.as_strided(index.shape, src.stride()) it = TensorIterator( outputs=[self_r], const_inputs=[src_r, index], check_mem_overlap=False, check_all_same_dtype=False, resize_outputs=False, ) except RuntimeError: return None if it.ndim != 2: return None elem = self.element_size() idx_elem = index.element_size() s_self, s_src, s_idx = it.strides(0), it.strides(1), it.strides(2) if ( s_idx[0] == 0 and s_idx[1] == idx_elem and s_self[0] == elem and s_src[0] == elem and s_self[1] == 0 ): return it return None def _alignment_contract_ok( dst: torch.Tensor, src: torch.Tensor, dim: int, *, dst_ptr_align: int ) -> bool: """Alignment of effective ptrs and outer row strides. The src side is always read via 16B-wide vector ops (LDG.128 / tile-mode TMA load), so its data_ptr and per-row stride must be 16B-aligned. The dst side's requirement is path-specific: - TMA: ``cp.reduce.async.bulk`` writes the gmem operand via a TMA descriptor that requires 16B alignment of both base and stride. - vec-scatter: writes via scalar fp32 atomicAdd / f16x2 / bf16x2 paired atomics, which need natural alignment of the 4B atomic operand on both base address and per-row stride. The "row stride" is the stride along ``dim`` -- the kernel uses it to step between rows of the (M, N) views ``_prepare_kernel_inputs`` builds. When this returns False the cond rejects and the call falls through to aten, which has its own predicate and safe fallback to ``vectorized_scatter_add_kernel_launch`` / ``indexFunc{Small,Large}Index``. """ if src.data_ptr() % _SRC_ALIGN_BYTES: return False if dst.data_ptr() % dst_ptr_align: return False elem = dst.element_size() if (src.stride(dim) * elem) % _SRC_ALIGN_BYTES: return False if (dst.stride(dim) * elem) % dst_ptr_align: return False return True # --------------------------------------------------------------------------- # Per-kernel eligibility. ``it.shape[0]`` is the slice extent (N elements) # after coalescing; that's what the alignment / divisibility checks key on. # --------------------------------------------------------------------------- def _is_tma_supported( self: torch.Tensor, dim: int, index: torch.Tensor, src: torch.Tensor ) -> bool: if not _has_sm90_plus(): return False d = _normalize_dim(dim, self.ndim) it = _scatter_add_eligibility(self, d, index, src) if it is None: return False if not _alignment_contract_ok(self, src, d, dst_ptr_align=_TMA_DST_ALIGN_BYTES): return False # TMA chunk_bytes = min(row_bytes, 512); row_bytes must evenly divide # by chunk_bytes and chunk_bytes must be 16-byte aligned. from .tma_kernel import row_shape_supported return row_shape_supported(self.dtype, it.shape[0]) def _is_vec_scatter_supported( self: torch.Tensor, dim: int, index: torch.Tensor, src: torch.Tensor ) -> bool: # ``red.global.add.noftz.bf16x2`` requires sm_90+. fp16x2 (sm_70+) and # scalar fp32 atomicAdd (sm_60+) work everywhere we care about. if self.dtype is torch.bfloat16 and not _has_sm90_plus(): return False d = _normalize_dim(dim, self.ndim) it = _scatter_add_eligibility(self, d, index, src) if it is None: return False if not _alignment_contract_ok(self, src, d, dst_ptr_align=_VEC_DST_ALIGN_BYTES): return False # Each lane owns vec_elems consecutive elements; loop is either # fully in-bounds or fully skipped per lane. Only requirement is # N % vec_elems == 0 (fp32: %4, halves: %8). from .vec_scatter_kernel import vec_elems_for return it.shape[0] % vec_elems_for(self.dtype) == 0 # --------------------------------------------------------------------------- # Kernel-input prep. After eligibility the kernel host functions take 2D # (M, N) tensors and a 1D index. We synthesize those views from the user's # original operands using the TI-approved geometry. # --------------------------------------------------------------------------- def _prepare_kernel_inputs( self: torch.Tensor, d: int, index: torch.Tensor, src: torch.Tensor, it: TensorIterator, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Build (self_2d, index_1d, src_2d) for the kernel host functions. Caller has verified eligibility via ``_scatter_add_eligibility``. ``d`` is the already-normalized scatter axis. ``it.shape[0]`` is the coalesced slice extent (product of every non-scatter dim of ``self``). The 2D views preserve the user's original stride along ``d`` (used by the kernel for the row stride). """ N, M_src = it.shape[0], it.shape[1] M_self = self.shape[d] self_2d = self.as_strided((M_self, N), (self.stride(d), 1)) src_2d = src.as_strided((M_src, N), (src.stride(d), 1)) # Eligibility verified the iter has post-coalesce strides # ``s_idx = (0, idx_elem)`` on ``index``; that means ``index`` carries # M_src consecutive elements at offset 0 with element stride 1, padded # by stride-0 broadcast axes. A single ``as_strided`` collapses that # to the 1D view the kernels expect. index_1d = index.as_strided((M_src,), (1,)) if index_1d.dtype != torch.int64: index_1d = index_1d.to(torch.int64) return self_2d, index_1d, src_2d # --------------------------------------------------------------------------- # Cond / impl factories. ``_make_cond`` wraps a support check in the shared # dispatch boilerplate; ``_make_impls`` produces the (functional, .out, # in-place) triple for a given kernel. # --------------------------------------------------------------------------- def _make_cond( support_check, *, out_dst_ptr_align: int | None = None, ) -> _OpCondFn: """Build a cond function that runs ``support_check`` behind the shared boilerplate (cutedsl availability, non-deterministic mode, CUDA tensors, non-COW, dtype/shape match on ``out`` when present). ``out_dst_ptr_align`` is None for the functional / inplace variants. For the ``.out`` variant it is the path-specific destination alignment (TMA: 16, vec-scatter: 4) used to alignment-check ``out``, which the ``.out`` impl runs the kernel on instead of ``self``. ``support_check`` is ``_is_tma_supported`` or ``_is_vec_scatter_supported`` -- both take ``(self, dim, index, src)`` and rebuild the analysis iter internally. """ if out_dst_ptr_align is not None: def _cond(self, dim, index, src, *, out): if not _base_cond_ok(self, index, src, out): return False if out.dtype != self.dtype or out.shape != self.shape: return False if out.ndim == 0: return False # ``out`` is the destination of a functional .out call; it # plays the role of self in the iter analysis. if not support_check(out, dim, index, src): return False # The .out impl runs the kernel on ``out``, not ``self`` (see # _make_impls.out_impl). support_check above validates ``out``'s # layout via TI strides, but we still need to alignment-check it # explicitly here in case ``out`` is a distinct buffer with # different alignment from ``self``. d = _normalize_dim(dim, out.ndim) return _alignment_contract_ok(out, src, d, dst_ptr_align=out_dst_ptr_align) return _cond def _cond(self, dim, index, src, *args, **kwargs): if not _base_cond_ok(self, index, src): return False return support_check(self, dim, index, src) return _cond def _copy_if_distinct(out: torch.Tensor, self: torch.Tensor) -> None: """Seed ``out`` with ``self``'s values when they're separate buffers. Cond guarantees shape match, so no resize is needed.""" if out.data_ptr() != self.data_ptr(): out.copy_(self) def _make_impls(kernel_getter): """Build (functional, out, in-place) impl callables. ``kernel_getter`` lazily imports the ``*_scatter_add_into`` kernel host function (we defer to keep the DSL runtime out of registration). Cond has already verified eligibility; we rebuild the iter inside the impl to recover the post-coalesce geometry that ``_prepare_kernel_inputs`` needs. The build is pure C++ and cheap. """ def _run(dst: torch.Tensor, dim: int, index, src) -> torch.Tensor: d = _normalize_dim(dim, dst.ndim) it = _scatter_add_eligibility(dst, d, index, src) if it is None: raise RuntimeError( "scatter_add cutedsl: cond approved but iter rebuild failed" ) # Mirrors aten's `if (iter.numel() == 0) return;` in # ScatterGatherKernel.cu -- the eligibility check is purely # layout-shaped, so empty inputs satisfy it but have no work. if it.numel == 0: return dst dst_2d, index_1d, src_2d = _prepare_kernel_inputs(dst, d, index, src, it) kernel_getter()(dst_2d, index_1d, src_2d) return dst def impl(self, dim, index, src, *args, **kwargs): return _run(self.clone(), dim, index, src) def out_impl(self, dim, index, src, *, out): _copy_if_distinct(out, self) return _run(out, dim, index, src) def inplace_impl(self, dim, index, src, *args, **kwargs): return _run(self, dim, index, src) return impl, out_impl, inplace_impl def _tma_kernel(): from .tma_kernel import tma_scatter_add_into return tma_scatter_add_into def _vs_kernel(): from .vec_scatter_kernel import vec_scatter_add_into return vec_scatter_add_into _tma_impl, _tma_out_impl, _tma_inplace_impl = _make_impls(_tma_kernel) _vs_impl, _vs_out_impl, _vs_inplace_impl = _make_impls(_vs_kernel) _tma_cond = _make_cond(_is_tma_supported) _tma_out_cond = _make_cond(_is_tma_supported, out_dst_ptr_align=_TMA_DST_ALIGN_BYTES) _vs_cond = _make_cond(_is_vec_scatter_supported) _vs_out_cond = _make_cond( _is_vec_scatter_supported, out_dst_ptr_align=_VEC_DST_ALIGN_BYTES ) # --------------------------------------------------------------------------- # Registration # --------------------------------------------------------------------------- # Dispatch priority is first-registered-wins. TMA's cond is strictly # narrower than vec-scatter's (TMA only fires on sm_90+ with stricter D # alignment), so TMA goes first. _PATHS = ( # (cond, impl, out_cond, out_impl, inplace_impl) (_tma_cond, _tma_impl, _tma_out_cond, _tma_out_impl, _tma_inplace_impl), (_vs_cond, _vs_impl, _vs_out_cond, _vs_out_impl, _vs_inplace_impl), ) def _register_path( dispatch_key: str, cond: _OpCondFn, impl: _OpImplFn, out_cond: _OpCondFn, out_impl: _OpImplFn, inplace_impl: _OpImplFn, ) -> None: # scatter_add_ has its own dispatcher entry separate from scatter_add # and scatter_add.out; it's structured-delegated to .out below the # dispatcher layer, so overriding .out alone doesn't catch it. for op_symbol, cond_fn, impl_fn in ( ("scatter_add", cond, impl), ("scatter_add.out", out_cond, out_impl), ("scatter_add_", cond, inplace_impl), ): cu.register_op_override( "aten", op_symbol, dispatch_key, cond=cond_fn, impl=impl_fn, allow_multiple_override=True, ) def register_to_dispatch() -> None: for path in _PATHS: _register_path("CUDA", *path)