Add check_bounds option to speed up ct.load/store

Signed-off-by: Ziheng Deng <zihengd@nvidia.com>

Author: gunnersdeng

changelog.d/check-bounds-load-store.md

@@ -0,0 +1,5 @@
+- Added a ``check_bounds`` option to ``ct.load()``, ``ct.store()`` and the
+  ``TiledView.load()`` / ``TiledView.store()`` methods. It defaults to ``True``.
+  when set to ``False``, it declares all elements in the tile is guaranteed to
+  stay within the array bounds and skips the out-of-bounds check to speed up the
+  load/store. Setting it to ``False`` requires tileiras 13.4+.

src/cuda/tile/_ir/ops.py

@@ -873,12 +873,25 @@ def _uniform_tuple(val: Any, *, rank: int):
     return (val,) * rank
 
 
+def _check_bounds_to_inbounds(check_bounds: Var, rank: int) -> tuple[bool, ...]:
+    check = require_constant_bool(check_bounds)
+    inbounds = _uniform_tuple(not check, rank=rank)
+    if not check:
+        cur_version = Builder.get_current().ir_ctx.tileiras_version
+        if cur_version < BytecodeVersion.V_13_4:
+            raise TileUnsupportedFeatureError(
+                f"'check_bounds=False' requires tileiras {BytecodeVersion.V_13_4.as_string()}"
+                f" or later. Current version is {cur_version.as_string()}.")
+    return inbounds
+
+
 @dataclass(eq=False)
 class TileLoad(Operation, opcode="tile_load", memory_effect=MemoryEffect.LOAD):
     latency: Optional[int] = attribute()
     allow_tma: Optional[bool] = attribute()
     memory_order: MemoryOrder = attribute(default=MemoryOrder.WEAK)
     memory_scope: MemoryScope = attribute(default=MemoryScope.NONE)
+    inbounds: tuple[bool, ...] = attribute(default=())
     view: Var = operand()
     index: tuple[Var, ...] = operand()
     token: Optional[Var] = operand(default=None)
@@ -910,14 +923,15 @@ def generate_bytecode(self, ctx: BytecodeContext) -> tuple[bc.Value, bc.Value]:
             memory_ordering_semantics=memory_order_to_bytecode[self.memory_order],
             memory_scope=memory_scope_to_bytecode[self.memory_scope],
             optimization_hints=ctx.load_store_hints(self.latency, self.allow_tma),
-            inbounds=_uniform_tuple(False, rank=len(self.index)),
+            inbounds=self.inbounds or _uniform_tuple(False, rank=len(self.index)),
         )
         return res, res_token
 
 
 def _tile_load_impl_inner(array: Var, index_items: tuple[Var, ...], shape: Sequence[int],
                           order: Sequence[int], padding_mode: PaddingMode,
                           latency: Var, allow_tma: Var,
+                          inbounds: tuple[bool, ...] = (),
                           traversal_steps: Optional[tuple[int, ...]] = None,
                           memory_order: MemoryOrder = MemoryOrder.WEAK,
                           memory_scope: MemoryScope = MemoryScope.NONE) -> Var:
@@ -938,7 +952,7 @@ def _tile_load_impl_inner(array: Var, index_items: tuple[Var, ...], shape: Seque
     result, _token = add_operation_variadic(TileLoad, (res_ty, TokenTy()),
                                             view=view, index=index_items, latency=latency,
                                             allow_tma=allow_tma, memory_order=memory_order,
-                                            memory_scope=memory_scope)
+                                            memory_scope=memory_scope, inbounds=inbounds)
     return reshape(result, shape)
 
 
@@ -1009,7 +1023,7 @@ def raw_array_memory_store_offset_impl(self: Var, offset: Var, value: Var,
 
 @impl(ct.load)
 def tile_load_impl(array: Var, index: Var, shape: Var, order: Var,
-                   padding_mode: Var, latency: Var, allow_tma: Var,
+                   padding_mode: Var, check_bounds: Var, latency: Var, allow_tma: Var,
                    memory_order: Var, memory_scope: Var) -> Var:
     array_ty = require_array_type(array)
     index_ty = require_index_or_index_tuple_type(index)
@@ -1022,11 +1036,12 @@ def tile_load_impl(array: Var, index: Var, shape: Var, order: Var,
                                    allow_0d_shape=True)
     order = require_constant_axis_order(order, array_ty.ndim)
     padding_mode = require_constant_enum(padding_mode, PaddingMode)
+    inbounds = _check_bounds_to_inbounds(check_bounds, array_ty.ndim)
     mem_order = require_constant_enum(memory_order, MemoryOrder)
     mem_scope = require_constant_enum(memory_scope, MemoryScope)
     validate_memory_order_and_scope(mem_order, mem_scope, TileLoad)
     return _tile_load_impl_inner(array, index_items, shape, order, padding_mode, latency, allow_tma,
-                                 memory_order=mem_order, memory_scope=mem_scope)
+                                 inbounds=inbounds, memory_order=mem_order, memory_scope=mem_scope)
 
 
 @dataclass(eq=False)
@@ -1035,6 +1050,7 @@ class TileStore(Operation, opcode="tile_store", memory_effect=MemoryEffect.STORE
     allow_tma: Optional[bool] = attribute()
     memory_order: MemoryOrder = attribute(default=MemoryOrder.WEAK)
     memory_scope: MemoryScope = attribute(default=MemoryScope.NONE)
+    inbounds: tuple[bool, ...] = attribute(default=())
     view: Var = operand()
     index: tuple[Var, ...] = operand()
     tile: Var = operand()
@@ -1066,12 +1082,13 @@ def generate_bytecode(self, ctx: BytecodeContext) -> bc.Value:
             memory_ordering_semantics=memory_order_to_bytecode[self.memory_order],
             memory_scope=memory_scope_to_bytecode[self.memory_scope],
             optimization_hints=ctx.load_store_hints(self.latency, self.allow_tma),
-            inbounds=_uniform_tuple(False, rank=len(self.index))
+            inbounds=self.inbounds or _uniform_tuple(False, rank=len(self.index))
         )
 
 
 def _tile_store_impl_inner(array: Var, index_items: tuple[Var, ...], tile: Var,
                            order: Sequence[int], latency: Var, allow_tma: Var,
+                           inbounds: tuple[bool, ...] = (),
                            traversal_steps: Optional[tuple[int, ...]] = None,
                            memory_order: MemoryOrder = MemoryOrder.WEAK,
                            memory_scope: MemoryScope = MemoryScope.NONE):
@@ -1092,12 +1109,12 @@ def _tile_store_impl_inner(array: Var, index_items: tuple[Var, ...], tile: Var,
                                    traversal_steps)
     add_operation(TileStore, TokenTy(), view=view, index=index_items, tile=tile,
                   latency=latency, allow_tma=allow_tma, memory_order=memory_order,
-                  memory_scope=memory_scope)
+                  memory_scope=memory_scope, inbounds=inbounds)
 
 
 @impl(ct.store)
 def tile_store_impl(array: Var, index: Var, tile: Var, order: Var,
-                    latency: Var, allow_tma: Var,
+                    check_bounds: Var, latency: Var, allow_tma: Var,
                     memory_order: Var, memory_scope: Var):
     array_ty = require_array_type(array)
     index_ty = require_index_or_index_tuple_type(index)
@@ -1108,11 +1125,12 @@ def tile_store_impl(array: Var, index: Var, tile: Var, order: Var,
 
     tile = implicit_cast(tile, array_ty.dtype, "Stored tile is incompatible with array's dtype")
     order = require_constant_axis_order(order, array_ty.ndim)
+    inbounds = _check_bounds_to_inbounds(check_bounds, array_ty.ndim)
     mem_order = require_constant_enum(memory_order, MemoryOrder)
     mem_scope = require_constant_enum(memory_scope, MemoryScope)
     validate_memory_order_and_scope(mem_order, mem_scope, TileStore)
     _tile_store_impl_inner(array, index_items, tile, order, latency, allow_tma,
-                           memory_order=mem_order, memory_scope=mem_scope)
+                           inbounds=inbounds, memory_order=mem_order, memory_scope=mem_scope)
 
 
 @dataclass(eq=False)
@@ -3158,23 +3176,27 @@ def tiled_view_num_tiles(self: Var, axis: Var) -> Var:
 
 
 @impl(ct.TiledView.load)
-def tiled_view_load_impl(self: Var, index: Var, latency: Var, allow_tma: Var) -> Var:
+def tiled_view_load_impl(self: Var, index: Var, check_bounds: Var, latency: Var,
+                         allow_tma: Var) -> Var:
     view_ty = require_tiled_view_type(self)
     index_ty = require_index_or_index_tuple_type(index)
     index_items = index.get_aggregate().items if isinstance(index_ty, TupleTy) else (index,)
     if view_ty.ndim != len(index_items):
         raise TileTypeError(f"Index size {len(index_items)}"
                             f" does not match the tiled view rank {view_ty.ndim}")
 
+    inbounds = _check_bounds_to_inbounds(check_bounds, view_ty.ndim)
     [array] = self.get_aggregate().as_tuple()
     order = get_default_order(view_ty.ndim)
     return _tile_load_impl_inner(array, index_items, view_ty.tile_shape, order,
                                  view_ty.padding_mode, latency, allow_tma,
+                                 inbounds=inbounds,
                                  traversal_steps=view_ty.traversal_steps)
 
 
 @impl(ct.TiledView.store)
-def tiled_view_store_impl(self: Var, index: Var, tile: Var, latency: Var, allow_tma: Var):
+def tiled_view_store_impl(self: Var, index: Var, tile: Var, check_bounds: Var, latency: Var,
+                          allow_tma: Var):
     view_ty = require_tiled_view_type(self)
     index_ty = require_index_or_index_tuple_type(index)
     index_items = index.get_aggregate().items if isinstance(index_ty, TupleTy) else (index,)
@@ -3187,12 +3209,14 @@ def tiled_view_store_impl(self: Var, index: Var, tile: Var, latency: Var, allow_
         raise TileTypeError(f"Tile shape {tile_ty.shape} is not broadcastable"
                             f" to the tiled view's tile shape {view_ty.tile_shape}")
 
+    inbounds = _check_bounds_to_inbounds(check_bounds, view_ty.ndim)
     tile = broadcast_to(tile, view_ty.tile_shape)
     tile = implicit_cast(tile, view_ty.dtype,
                          "Stored tile is incompatible with tiled view's dtype")
     [array] = self.get_aggregate().as_tuple()
     order = get_default_order(view_ty.ndim)
     _tile_store_impl_inner(array, index_items, tile, order, latency, allow_tma,
+                           inbounds=inbounds,
                            traversal_steps=view_ty.traversal_steps)
 
 

src/cuda/tile/_stub.py

@@ -816,6 +816,7 @@ def traversal_steps(self) -> tuple[int, ...]:
 
     @stub
     def load(self, index: Shape, *,
+             check_bounds: Constant[bool] = True,
              latency: Optional[int] = None,
              allow_tma: Optional[bool] = None) -> Tile:
         """Loads a tile from the |tiled view| at the given tile `index`.
@@ -828,6 +829,11 @@ def load(self, index: Shape, *,
 
         Args:
             index (tuple[int,...]): An index in the |tiled view|'s tile space.
+            check_bounds (const bool): Whether to bounds-check the tile against the view
+                boundaries. When ``False``, the tile is assumed to stay fully within bounds and the
+                out-of-bounds check is skipped, producing a faster load; violating this assumption
+                is undefined behavior. Defaults to ``True``. Setting it to ``False`` requires
+                tileiras 13.4+.
             latency (const int): A hint indicating how heavy DRAM traffic will be. It shall be an
                 integer between 1 (low) and 10 (high). By default, the compiler will infer the
                 latency.
@@ -858,6 +864,7 @@ def kernel(x):
 
     @stub
     def store(self, index: Shape, tile: Tile, *,
+              check_bounds: Constant[bool] = True,
               latency: Optional[int] = None,
               allow_tma: Optional[bool] = None) -> None:
         """Stores a `tile` into the |tiled view| at the given tile `index`.
@@ -872,6 +879,11 @@ def store(self, index: Shape, tile: Tile, *,
         Args:
             index (tuple[int,...]): An index in the |tiled view|'s tile space.
             tile (Tile): The tile to store.
+            check_bounds (const bool): Whether to bounds-check the tile against the view
+                boundaries. When ``False``, the tile is assumed to stay fully within bounds and the
+                out-of-bounds check is skipped, producing a faster store; violating this assumption
+                is undefined behavior. Defaults to ``True``. Setting it to ``False`` requires
+                tileiras 13.4+.
             latency (const int): A hint indicating how heavy DRAM traffic will be. It shall be an
                 integer between 1 (low) and 10 (high). By default, the compiler will infer the
                 latency.
@@ -1255,6 +1267,7 @@ def load(array: Array, /,
          shape: Constant[Shape], *,
          order: Constant[Order] = "C",
          padding_mode: PaddingMode = PaddingMode.UNDETERMINED,
+         check_bounds: Constant[bool] = True,
          latency: Optional[int] = None,
          allow_tma: Optional[bool] = None,
          memory_order: MemoryOrder = MemoryOrder.WEAK,
@@ -1302,6 +1315,11 @@ def load(array: Array, /,
 
         padding_mode (PaddingMode): The value used to pad the tile when it extends beyond the array
             boundaries. By default, the padding value is undetermined.
+        check_bounds (const bool): Whether to bounds-check the tile against the array boundaries.
+            When ``False``, the tile is assumed to stay fully within bounds and the out-of-bounds
+            check is skipped, producing a faster load; violating this assumption is undefined
+            behavior. Since no out-of-bound elements can occur in this case, ``padding_mode`` is
+            ignored. Defaults to ``True``. Setting it to ``False`` requires tileiras 13.4+.
         latency (const int): A hint indicating how heavy DRAM traffic will be. It shall be an
             integer between 1 (low) and 10 (high). By default, the compiler will infer the latency.
         allow_tma (const bool): If False, the load will not use TMA. By default, TMA is allowed.
@@ -1408,6 +1426,7 @@ def store(array: Array, /,
           index: Shape,
           tile: TileOrScalar, *,
           order: Constant[Order] = "C",
+          check_bounds: Constant[bool] = True,
           latency: Optional[int] = None,
           allow_tma: Optional[bool] = None,
           memory_order: MemoryOrder = MemoryOrder.WEAK,
@@ -1437,6 +1456,10 @@ def store(array: Array, /,
         tile (Tile): The |tile| to store. The rank of the tile must match rank of the array,
             unless it is a scalar or 0d tile.
         order ("C" or "F", or tuple[const int,...]): Order of axis mapping. See :py:func:`load`.
+        check_bounds (const bool): Whether to bounds-check the tile against the array boundaries.
+            When ``False``, the tile is assumed to stay fully within bounds and the out-of-bounds
+            check is skipped, producing a faster store; violating this assumption is undefined
+            behavior. Defaults to ``True``. Setting it to ``False`` requires tileiras 13.4+.
         latency (int, optional): A hint indicating how heavy DRAM traffic will be. It shall be an
             integer between 1 (low) and 10 (high). By default, the compiler will infer the latency.
         allow_tma (bool, optional): If False, the store will not use TMA.

test/test_load_store.py

@@ -7,8 +7,9 @@
 import pytest
 
 from math import ceil
-from conftest import float_dtypes, bool_dtypes, get_tileiras_version, int_dtypes, dtype_id
-from cuda.tile import TileTypeError
+from conftest import (float_dtypes, bool_dtypes, get_tileiras_version, int_dtypes, dtype_id,
+                      requires_tileiras)
+from cuda.tile import TileTypeError, TileUnsupportedFeatureError
 from cuda.tile._bytecode.version import BytecodeVersion
 from cuda.tile._ir.cast_ops import _is_implicit_cast_ok
 from cuda.tile._ir.typing_support import to_dtype
@@ -214,3 +215,51 @@ def kern(x):
                        match="Axis order must be a permutation, but axis 1 is used at least twice"):
         x = torch.zeros((64, 64), device="cuda")
         ct.launch(torch.cuda.current_stream(), (1,), kern, (x,))
+
+
+@ct.kernel
+def copy_2d_no_check_bounds(x, y, TILE_X: ct.Constant[int], TILE_Y: ct.Constant[int]):
+    bidx = ct.bid(0)
+    bidy = ct.bid(1)
+    tx = ct.load(x, index=(bidx, bidy), shape=(TILE_X, TILE_Y), check_bounds=False)
+    ct.store(y, index=(bidx, bidy), tile=tx, check_bounds=False)
+
+
+@pytest.mark.use_mlir
+@requires_tileiras(BytecodeVersion.V_13_4)
+def test_load_store_check_bounds():
+    # check_bounds=False lowers to an all-true `inbounds` attribute on every dimension.
+    shape = (64, 64)
+    tile = (32, 32)
+    x = make_tensor(shape, dtype=torch.float32, device="cuda")
+    y = torch.zeros_like(x)
+    grid = (shape[0] // tile[0], shape[1] // tile[1], 1)
+    bytecode = get_bytecode(copy_2d_no_check_bounds, (x, y, tile[0], tile[1]))
+    wildcard = "{{.*}}"
+    filecheck(bytecode, "\n".join([
+        f"// CHECK: load_view_tko{wildcard}inbounds = [true, true]",
+        f"// CHECK: store_view_tko{wildcard}inbounds = [true, true]",
+    ]))
+    ct.launch(torch.cuda.current_stream(), grid, copy_2d_no_check_bounds, (x, y, tile[0], tile[1]))
+    assert_equal(y, x)
+
+
+@pytest.mark.use_mlir
+def test_load_store_check_bounds_default():
+    # The default check_bounds=True emits no `inbounds` attribute.
+    x = make_tensor((32,), dtype=torch.float16, device="cuda")
+    y = torch.zeros_like(x)
+    bytecode = get_bytecode(array_copy_1d, (x, y, 16))
+    filecheck(bytecode, "\n".join([
+        "// CHECK: load_view_tko",
+        "// CHECK-NOT: inbounds",
+    ]))
+
+
+@pytest.mark.skipif(get_tileiras_version() >= BytecodeVersion.V_13_4,
+                    reason="check_bounds=False is supported on tileiras 13.4+")
+def test_check_bounds_requires_13_4():
+    x = make_tensor((64, 64), dtype=torch.float16, device="cuda")
+    y = torch.zeros_like(x)
+    with pytest.raises(TileUnsupportedFeatureError, match="check_bounds=False.*requires tileiras"):
+        get_bytecode(copy_2d_no_check_bounds, (x, y, 32, 32))