Support GatherScatterView via advanced indexing

Signed-off-by: Qiqi Xiao <qiqix@nvidia.com>

Author: xiaoqiqi177

changelog.d/load-store-advanced.md

@@ -0,0 +1,4 @@
+<!--- SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!--- SPDX-License-Identifier: Apache-2.0 -->
+
+- New `ct.load_advanced(array, indices)` and `ct.store_advanced(array, indices, tile)` for gathering/scattering along one dimension from/to a 2D or higher-rank array. A 1D integer `Tile` selects the sparse dimension; `ct.Slice(start, length)` selects a contiguous range along each dense dimension.

docs/source/data/slice.rst

@@ -0,0 +1,16 @@
+.. SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+..
+.. SPDX-License-Identifier: Apache-2.0
+
+.. currentmodule:: cuda.tile
+
+.. _data-slice-cuda-tile-slice:
+
+cuda.tile.Slice
+===============
+
+.. autoclass:: Slice
+   :members:
+   :undoc-members:
+   :special-members:
+   :exclude-members: __annotations__, __dict__, __module__, __weakref__

docs/source/operations.rst

@@ -22,6 +22,8 @@ Load/Store
    num_tiles
    load
    store
+   load_advanced
+   store_advanced
    gather
    scatter
 
@@ -234,12 +236,14 @@ Classes
 
    Array
    TiledView
+   Slice
 
 .. toctree::
    :hidden:
 
    data/array
    data/tiled_view
+   data/slice
 
 
 .. _operations-enums:

src/cuda/tile/__init__.py

@@ -64,6 +64,7 @@
     ListAnnotation,
     Scalar,
     ScalarInt64,
+    Slice,
     Tile,
     TiledView,
 
@@ -114,6 +115,7 @@
     less,
     less_equal,
     load,
+    load_advanced,
     log,
     log2,
     matmul,
@@ -148,6 +150,7 @@
     static_eval,
     static_iter,
     store,
+    store_advanced,
     sub,
     sum,
     tan,
@@ -222,6 +225,7 @@
     "ListAnnotation",
     "Scalar",
     "ScalarInt64",
+    "Slice",
     "Tile",
     "TiledView",
 
@@ -272,6 +276,7 @@
     "less",
     "less_equal",
     "load",
+    "load_advanced",
     "log",
     "log2",
     "matmul",
@@ -306,6 +311,7 @@
     "static_eval",
     "static_iter",
     "store",
+    "store_advanced",
     "sub",
     "sum",
     "tan",

src/cuda/tile/_ir/ir.py

@@ -309,6 +309,15 @@ def as_tuple(self) -> tuple["Var", ...]:
         return (self.array,)
 
 
+@dataclass
+class IndexSliceValue(AggregateValue):
+    start: Var
+    length: Var
+
+    def as_tuple(self) -> tuple[Var, ...]:
+        return (self.start, self.length)
+
+
 @dataclass
 class RawArrayMemoryValue(AggregateValue):
     base_ptr: Var

src/cuda/tile/_ir/ops.py

@@ -27,7 +27,8 @@
     enter_nested_block, nested_block, PhiState, LoopVarState,
     TupleValue, make_aggregate, RangeValue, BoundMethodValue, ArrayValue, ConstantState,
     ListValue, TiledViewValue, ClosureValue, MemoryEffect, attribute, operand,
-    BlockRestriction, FormattedStringValue, RawArrayMemoryValue, DataclassValue, DataclassInfo
+    BlockRestriction, FormattedStringValue, RawArrayMemoryValue, DataclassValue, DataclassInfo,
+    IndexSliceValue
 )
 from .type import PointerTy
 from . import hir, hir_stubs
@@ -59,11 +60,12 @@
     typeof_pyval, dtype_registry, loose_type_of_pyval, get_constant_value, get_dataclass_info,
 )
 from .type import (
-    PartitionViewTy, StridedViewTy, TupleTy, TileTy, NoneType, BoundMethodTy, ArrayTy,
+    PartitionViewTy, StridedViewTy, GatherScatterViewTy, TupleTy, TileTy, NoneType,
+    BoundMethodTy, ArrayTy,
     ListTy, make_tile_ty, SliceType, DTypeConstructor, RangeIterType, Type,
     NONE, ModuleTy, TypeTy, LooselyTypedScalar, DTypeSpec, StringTy, InvalidType,
     ClosureTy, LiveCapturedScope, TokenTy, TiledViewTy, FormattedStringTy,
-    StringFormat, FormattedPiece, RawArrayMemoryTy, DataclassTy
+    StringFormat, FormattedPiece, RawArrayMemoryTy, DataclassTy, IndexSliceTy
 )
 from cuda.tile._datatype import (
     DType, is_integral, is_float, is_signed, is_boolean,
@@ -2375,6 +2377,29 @@ def _materialize_tiled_view(array: Var,
     return _make_partition_view(array, tile_shape, order, padding_mode)
 
 
+@dataclass(eq=False)
+class MakeGatherScatterView(Operation, opcode="make_gather_scatter_view"):
+    array: Var = operand()
+
+    @override
+    def generate_bytecode(self, ctx: BytecodeContext) -> bc.Value:
+        gs_view_ty = self.result_var.get_type()
+        return bc.encode_MakeGatherScatterViewOp(ctx.builder,
+                                                 typeid(ctx.type_table, gs_view_ty),
+                                                 ctx.get_value(self.array))
+
+
+def make_gather_scatter_view(array: Var, tile_shape: Sequence[int],
+                             sparse_dim: int,
+                             padding_mode: PaddingMode) -> Var:
+    array_ty = array.get_type()
+    assert isinstance(array_ty, ArrayTy)
+    view_ty = GatherScatterViewTy(array_ty, tuple(tile_shape), sparse_dim, padding_mode)
+    ret = add_operation(MakeGatherScatterView, view_ty, array=array)
+    ret.set_aggregate(array.get_aggregate())
+    return ret
+
+
 @dataclass(eq=False)
 class TileLoad(Operation, opcode="tile_load", memory_effect=MemoryEffect.LOAD):
     latency: Optional[int] = attribute()
@@ -4826,6 +4851,128 @@ def tiled_view_atomic_rmw_impl(int_mode: Optional[AtomicRMWMode],
                   view=view, index=index_items, update=update)
 
 
+@impl(ct.Slice)
+def slice_index_constructor_impl(start: Var, length: Var) -> Var:
+    start_ty = require_signed_integer_0d_tile_type(start)
+    length_ty = require_signed_integer_0d_tile_type(length)
+    res_type = IndexSliceTy(start_ty, length_ty)
+    res_loose_type = IndexSliceTy(start.get_loose_type(), length.get_loose_type())
+    return make_aggregate(IndexSliceValue(start, length), res_type, res_loose_type)
+
+
+def _parse_advanced_index(indices: Var, ndim: int) -> tuple[int, tuple[int, ...], tuple[Var, ...]]:
+    """Unpack, classify, validate, and build the gather scatter view index.
+
+    Returns (sparse_dim, tile_shape, gs_index).
+    """
+    require_tuple_type(indices)
+    items = list(indices.get_aggregate().items)
+    if len(items) != ndim:
+        raise TileTypeError(
+            f"load_advanced/store_advanced index length {len(items)} does not "
+            f"match array rank {ndim}")
+
+    sparse_dims: list[int] = []
+    tile_shape: list[int] = []
+    gs_index: list[Var] = []
+
+    for dim, item in enumerate(items):
+        item_ty = item.get_type()
+        if isinstance(item_ty, TileTy):
+            if item_ty.ndim != 1:
+                raise TileTypeError(
+                    f"Sparse index at dim {dim} must be a 1D integer tile, "
+                    f"got {item_ty.ndim}D")
+            if not is_integral(item_ty.dtype):
+                raise TileTypeError(
+                    f"Sparse index at dim {dim} must be an integer tile, "
+                    f"got dtype {item_ty.dtype}")
+            sparse_dims.append(dim)
+            tile_shape.append(item_ty.shape[0])
+            gs_index.append(item)
+        elif isinstance(item_ty, IndexSliceTy):
+            length_var = item.get_aggregate().length
+            if not length_var.is_constant():
+                raise TileTypeError(
+                    f"ct.Slice length at dim {dim} must be a compile-time constant "
+                    f"in load_advanced/store_advanced")
+            length_val = length_var.get_constant()
+            if not isinstance(length_val, int) or length_val <= 0:
+                raise TileTypeError(
+                    f"ct.Slice length at dim {dim} must be a positive integer, got {length_val}")
+            tile_shape.append(length_val)
+            gs_index.append(item.get_aggregate().start)
+        else:
+            raise TileTypeError(
+                f"load_advanced/store_advanced index at dim {dim} must be a "
+                f"1D integer Tile (sparse dim) or ct.Slice(start, length) "
+                f"(dense dim), got type {item_ty}")
+
+    if len(sparse_dims) == 0:
+        raise TileTypeError(
+            "load_advanced/store_advanced: exactly one index must be a 1D "
+            "integer Tile (the sparse dim); none found")
+    if len(sparse_dims) > 1:
+        raise TileTypeError(
+            f"load_advanced/store_advanced: exactly one index must be a 1D "
+            f"integer Tile (the sparse dim); found {len(sparse_dims)} at "
+            f"dims {sparse_dims}")
+
+    for dim, n in enumerate(tile_shape):
+        if not _is_power_of_2(n):
+            raise TileTypeError(
+                f"Index at dim {dim} has size {n}; must be a power of two")
+
+    return sparse_dims[0], tuple(tile_shape), tuple(gs_index)
+
+
+@impl(ct.load_advanced, min_version=BytecodeVersion.V_13_3)
+def load_advanced_impl(array: Var, indices: Var, padding_mode: Var,
+                       latency: Var, allow_tma: Var) -> Var:
+    array_ty = require_array_type(array)
+    if array_ty.ndim < 2:
+        raise TileTypeError(
+            "load_advanced requires a 2D or higher-rank array; "
+            "use ct.gather() for 1D arrays")
+    sparse_dim, tile_shape, gs_index = _parse_advanced_index(indices, array_ty.ndim)
+    padding_mode_val = require_constant_enum(padding_mode, PaddingMode)
+    latency_val = require_optional_constant_int(latency)
+    allow_tma_val = require_optional_constant_bool(allow_tma)
+    _check_load_store_hints(latency_val, allow_tma_val)
+
+    view = make_gather_scatter_view(array, tile_shape, sparse_dim, padding_mode_val)
+    result, _token = add_operation(TileLoad, (make_tile_ty(array_ty.dtype, tile_shape), TokenTy()),
+                                   view=view, index=gs_index,
+                                   latency=latency_val, allow_tma=allow_tma_val)
+    return result
+
+
+@impl(ct.store_advanced, min_version=BytecodeVersion.V_13_3)
+def store_advanced_impl(array: Var, indices: Var, tile: Var,
+                        latency: Var, allow_tma: Var):
+    array_ty = require_array_type(array)
+    if array_ty.ndim < 2:
+        raise TileTypeError(
+            "store_advanced requires a 2D or higher-rank array; "
+            "use ct.scatter() for 1D arrays")
+    sparse_dim, tile_shape, gs_index = _parse_advanced_index(indices, array_ty.ndim)
+    tile_ty = require_tile_type(tile)
+    if tile_ty.shape != tile_shape:
+        raise TileTypeError(
+            f"Tile shape {tile_ty.shape} does not match the shape implied by "
+            f"indices {tile_shape}")
+    tile = _implicit_cast(tile, array_ty.dtype,
+                          "Stored tile dtype is incompatible with array dtype")
+    latency_val = require_optional_constant_int(latency)
+    allow_tma_val = require_optional_constant_bool(allow_tma)
+    _check_load_store_hints(latency_val, allow_tma_val)
+
+    view = make_gather_scatter_view(array, tile_shape, sparse_dim, PaddingMode.UNDETERMINED)
+    [_token] = add_operation(TileStore, (TokenTy(),),
+                             view=view, index=gs_index, tile=tile,
+                             latency=latency_val, allow_tma=allow_tma_val)
+
+
 def store_var(local_idx: int, value: Var, loc: Loc | None = None):
     scope = Scope.get_current()
     new_var = scope.local.redefine(local_idx, loc or Builder.get_current().loc)

src/cuda/tile/_ir/type.py

@@ -450,6 +450,58 @@ def __str__(self):
                 f"padding_mode={self.padding_mode}]")
 
 
+# ============== GatherScatterView Type ===============
+
+
+@dataclass(frozen=True)
+class GatherScatterViewTy(Type):
+    array_ty: ArrayTy
+    tile_shape: tuple[int, ...]
+    sparse_dim: int
+    padding_mode: PaddingMode
+
+    def is_aggregate(self) -> bool:
+        return True
+
+    def aggregate_item_types(self) -> tuple["Type", ...]:
+        return self.array_ty.aggregate_item_types()
+
+    def make_aggregate_value(self, items: tuple["Var", ...]) -> "AggregateValue":
+        return self.array_ty.make_aggregate_value(items)
+
+    @property
+    def dtype(self):
+        return self.array_ty.dtype
+
+    def __str__(self):
+        return (f"GatherScatterView[{self.array_ty},tile_shape={self.tile_shape},"
+                f"sparse_dim={self.sparse_dim},padding_mode={self.padding_mode}]")
+
+
+# ============== IndexSlice Type ===============
+
+
+@dataclass(frozen=True)
+class IndexSliceTy(Type):
+    """Type of a ct.Slice(start, length)."""
+    start_ty: "Type"
+    length_ty: "Type"
+
+    def is_aggregate(self) -> bool:
+        return True
+
+    def aggregate_item_types(self) -> tuple["Type", ...]:
+        return (self.start_ty, self.length_ty)
+
+    def make_aggregate_value(self, items: tuple["Var", ...]) -> "AggregateValue":
+        from .ir import IndexSliceValue
+        assert len(items) == 2
+        return IndexSliceValue(items[0], items[1])
+
+    def __str__(self) -> str:
+        return f"IndexSlice[start_ty={self.start_ty}, length_ty={self.length_ty}]"
+
+
 # ============== TiledView Type ===============
 
 @dataclass(frozen=True)

src/cuda/tile/_ir/typing_support.py

@@ -118,7 +118,7 @@ def get_signature(f) -> inspect.Signature:
 
 
 def is_supported_builtin_func(x: Any) -> bool:
-    return _safe_get(BUILTIN_FUNCS, x) is not None
+    return _safe_get(BUILTIN_FUNCS, x) is not None or getattr(x, '_cutile_is_builtin', False)
 
 
 def typeof_pyval(val) -> Type:

src/cuda/tile/_ir2bytecode.py

@@ -20,8 +20,8 @@
     get_default_rounding_mode,
 )
 from cuda.tile._ir.type import (
-    PartitionViewTy, StridedViewTy, Type, TileTy, PointerTy, TokenTy, TupleTy, ArrayTy,
-    size_to_bytecode,
+    PartitionViewTy, StridedViewTy, GatherScatterViewTy, Type, TileTy, PointerTy, TokenTy,
+    TupleTy, ArrayTy, size_to_bytecode,
 )
 
 
@@ -61,6 +61,11 @@ def typeid(tt: bc.TypeTable, ty: Type) -> bc.TypeId:
                                    ty.order, padding_value)
         else:
             return tt.partition_view(ty.tile_shape, tv_id, ty.order, padding_value)
+    elif isinstance(ty, GatherScatterViewTy):
+        padding_value = padding_mode_to_bytecode[ty.padding_mode]
+        assert isinstance(ty.array_ty, ArrayTy)
+        tv_id = tensor_view_typeid(tt, ty.array_ty)
+        return tt.gather_scatter_view(ty.tile_shape, tv_id, ty.sparse_dim, padding_value)
     else:
         raise NotImplementedError(f"Lowering type '{ty}' is not supported")