[KernelGen][thead] Add unsafe_masked_index_put_accumulate vendor specialization (#206)

Co-authored-by: Dongxu-H <dxhan@baai.ac.cn>

Author: chx7514

src/flag_gems/runtime/backend/_thead/ops/__init__.py

@@ -30,10 +30,12 @@
 from .special_shifted_chebyshev_polynomial_w import (
     special_shifted_chebyshev_polynomial_w,
 )
+from .unsafe_masked_index_put_accumulate import _unsafe_masked_index_put_accumulate
 
 __all__ = [
     "_conv_depthwise2d",
     "_thnn_fused_lstm_cell_backward_impl",
+    "_unsafe_masked_index_put_accumulate",
     "adaptive_max_pool3d_backward",
     "broadcast_tensors",
     "broadcast_to",

src/flag_gems/runtime/backend/_thead/ops/unsafe_masked_index_put_accumulate.py

@@ -0,0 +1,258 @@
+# Copyright 2026, The FlagOS Contributors.
+#
+# 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.
+#
+# Generated by KernelGen: https://github.com/flagos-ai/KernelGen
+
+import logging
+
+import triton
+import triton.language as tl
+
+from flag_gems.utils import libentry
+
+logger = logging.getLogger(__name__)
+
+# Per-dimension index tensors are gathered into the kernel through a fixed set
+# of pointer arguments. Unused slots are padded and never read because RANK is
+# a compile-time constant that guards every access.
+_MAX_RANK = 8
+
+
+@libentry()
+@triton.jit
+def _masked_scatter_accumulate_flat_kernel(
+    out_ptr,
+    mask_ptr,
+    indices_ptr,
+    values_ptr,
+    n_elements,
+    out_numel,
+    BLOCK_SIZE: tl.constexpr,
+):
+    # Fast path: indices are already a single flat linear-index tensor, so the
+    # kernel is a coalesced load + masked atomic-add with no index math.
+    pid = tl.program_id(axis=0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    lane_mask = offsets < n_elements
+
+    mask_val = tl.load(mask_ptr + offsets, mask=lane_mask, other=0)
+    idx = tl.load(indices_ptr + offsets, mask=lane_mask, other=0).to(tl.int64)
+    vals = tl.load(values_ptr + offsets, mask=lane_mask, other=0)
+
+    index_valid = (idx >= 0) & (idx < out_numel)
+    final_mask = lane_mask & (mask_val != 0) & index_valid
+
+    tl.atomic_add(out_ptr + idx, vals, mask=final_mask, sem="relaxed")
+
+
+@libentry()
+@triton.jit
+def _masked_scatter_accumulate_nd_kernel(
+    out_ptr,
+    mask_ptr,
+    values_ptr,
+    idx_ptr0,
+    idx_ptr1,
+    idx_ptr2,
+    idx_ptr3,
+    idx_ptr4,
+    idx_ptr5,
+    idx_ptr6,
+    idx_ptr7,
+    stride0,
+    stride1,
+    stride2,
+    stride3,
+    stride4,
+    stride5,
+    stride6,
+    stride7,
+    n_elements,
+    out_numel,
+    RANK: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    # Fused path: per-dimension index tensors are combined into a flat linear
+    # index in-kernel, then accumulated. This removes the several separate
+    # torch launches an equivalent host-side flatten would cost on the PPU.
+    pid = tl.program_id(axis=0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    lane_mask = offsets < n_elements
+
+    idx = tl.load(idx_ptr0 + offsets, mask=lane_mask, other=0).to(tl.int64) * stride0
+    if RANK > 1:
+        idx += (
+            tl.load(idx_ptr1 + offsets, mask=lane_mask, other=0).to(tl.int64) * stride1
+        )
+    if RANK > 2:
+        idx += (
+            tl.load(idx_ptr2 + offsets, mask=lane_mask, other=0).to(tl.int64) * stride2
+        )
+    if RANK > 3:
+        idx += (
+            tl.load(idx_ptr3 + offsets, mask=lane_mask, other=0).to(tl.int64) * stride3
+        )
+    if RANK > 4:
+        idx += (
+            tl.load(idx_ptr4 + offsets, mask=lane_mask, other=0).to(tl.int64) * stride4
+        )
+    if RANK > 5:
+        idx += (
+            tl.load(idx_ptr5 + offsets, mask=lane_mask, other=0).to(tl.int64) * stride5
+        )
+    if RANK > 6:
+        idx += (
+            tl.load(idx_ptr6 + offsets, mask=lane_mask, other=0).to(tl.int64) * stride6
+        )
+    if RANK > 7:
+        idx += (
+            tl.load(idx_ptr7 + offsets, mask=lane_mask, other=0).to(tl.int64) * stride7
+        )
+
+    mask_val = tl.load(mask_ptr + offsets, mask=lane_mask, other=0)
+    vals = tl.load(values_ptr + offsets, mask=lane_mask, other=0)
+
+    index_valid = (idx >= 0) & (idx < out_numel)
+    final_mask = lane_mask & (mask_val != 0) & index_valid
+
+    tl.atomic_add(out_ptr + idx, vals, mask=final_mask, sem="relaxed")
+
+
+def _unsafe_masked_index_put_accumulate(inp, mask, indices, values):
+    logger.debug("GEMS_HEAD UNSAFE_MASKED_INDEX_PUT_ACCUMULATE")
+
+    # Normalize the indices argument. torch passes Tensor?[] (a list/tuple of
+    # per-dimension index tensors); the generic layer's internal calls pass a
+    # single flat linear-index tensor.
+    per_dim = None
+    flat_indices = None
+    if isinstance(indices, (list, tuple)):
+        if len(indices) == 0 or indices[0] is None:
+            raise ValueError("Empty indices list")
+        if len(indices) == 1:
+            flat_indices = indices[0]
+        else:
+            per_dim = list(indices)
+    else:
+        flat_indices = indices
+
+    ref_shape = mask.shape
+    assert mask.shape == values.shape, (
+        f"mask and values must have same shape, got {mask.shape} " f"and {values.shape}"
+    )
+
+    # Co-locate everything on the value tensor's device.
+    dev = values.device
+    if inp.device != dev:
+        inp = inp.to(dev)
+    if mask.device != dev:
+        mask = mask.to(dev)
+    if values.device != dev:
+        values = values.to(dev)
+
+    # Work on a contiguous copy so the flat linear index equals the memory
+    # offset and the input is never mutated (out-of-place semantics).
+    out = inp.contiguous()
+    if out.data_ptr() == inp.data_ptr():
+        out = out.clone()
+
+    n_elements = mask.numel()
+    out_numel = out.numel()
+    if n_elements == 0:
+        return out
+
+    mask_flat = mask.contiguous().reshape(-1)
+    values_flat = values.contiguous().reshape(-1)
+    out_flat = out.reshape(-1)
+
+    # BLOCK_SIZE=256 / num_warps=1: tuned for the thead PPU scatter workload —
+    # low launch overhead on the tiny shapes that dominate this op while giving
+    # enough parallelism for the atomic accumulate on larger inputs.
+    BLOCK_SIZE = 256
+    grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
+
+    if per_dim is not None:
+        rank = len(per_dim)
+        if rank > _MAX_RANK:
+            raise ValueError(f"rank {rank} exceeds supported maximum {_MAX_RANK}")
+        # Row-major logical strides of the input: stride[i] = prod(shape[i+1:]).
+        logical_strides = [1] * rank
+        acc = 1
+        for i in range(rank - 1, -1, -1):
+            logical_strides[i] = acc
+            acc *= inp.shape[i]
+
+        idx_ptrs = []
+        for t in per_dim:
+            assert t.shape == ref_shape, (
+                f"mask and indices must have same shape, got {ref_shape} "
+                f"and {t.shape}"
+            )
+            it = t if t.device == dev else t.to(dev)
+            idx_ptrs.append(it.contiguous().reshape(-1))
+        # Pad pointer / stride slots up to _MAX_RANK; padded slots are guarded
+        # out by the RANK constexpr and never dereferenced.
+        pad_ptr = idx_ptrs[0]
+        while len(idx_ptrs) < _MAX_RANK:
+            idx_ptrs.append(pad_ptr)
+        strides = logical_strides + [0] * (_MAX_RANK - rank)
+
+        _masked_scatter_accumulate_nd_kernel[grid](
+            out_flat,
+            mask_flat,
+            values_flat,
+            idx_ptrs[0],
+            idx_ptrs[1],
+            idx_ptrs[2],
+            idx_ptrs[3],
+            idx_ptrs[4],
+            idx_ptrs[5],
+            idx_ptrs[6],
+            idx_ptrs[7],
+            strides[0],
+            strides[1],
+            strides[2],
+            strides[3],
+            strides[4],
+            strides[5],
+            strides[6],
+            strides[7],
+            n_elements,
+            out_numel,
+            RANK=rank,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=1,
+        )
+        return out
+
+    # Single flat-index path.
+    assert flat_indices.shape == ref_shape, (
+        f"mask and indices must have same shape, got {ref_shape} "
+        f"and {flat_indices.shape}"
+    )
+    if flat_indices.device != dev:
+        flat_indices = flat_indices.to(dev)
+    indices_flat = flat_indices.contiguous().reshape(-1)
+
+    _masked_scatter_accumulate_flat_kernel[grid](
+        out_flat,
+        mask_flat,
+        indices_flat,
+        values_flat,
+        n_elements,
+        out_numel,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=1,
+    )
+    return out