Add RawArrayMemory() to support array direct memory access by element offset

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

Author: gunnersdeng

changelog.d/raw-array-memory.md

@@ -0,0 +1 @@
+- Added `Array.get_raw_memory()` returning a `RawArrayMemory` object that supports `load_offset(offset)` and `store_offset(offset, value)` for direct memory access by element offset (no shape/stride index calculation).

src/cuda/tile/_ir/ir.py

@@ -288,6 +288,14 @@ def as_tuple(self) -> tuple["Var", ...]:
         return (self.array,)
 
 
+@dataclass
+class RawArrayMemoryValue(AggregateValue):
+    base_ptr: Var
+
+    def as_tuple(self) -> tuple[Var, ...]:
+        return (self.base_ptr,)
+
+
 @dataclass
 class ListValue(AggregateValue):
     base_ptr: Var

src/cuda/tile/_ir/op_impl.py

@@ -20,7 +20,7 @@
 from .ir import Var, TupleValue, Builder
 from .type import TiledViewTy, TupleTy, TileTy, DTypeSpec, EnumTy, StringTy, ArrayTy, SliceType, \
     ListTy, LooselyTypedScalar, RangeIterType, FunctionTy, ClosureTy, BoundMethodTy, \
-    DTypeConstructor, Type
+    DTypeConstructor, Type, RawArrayMemoryTy
 
 
 def _verify_params_match(stub_sig: inspect.Signature, func_sig: inspect.Signature):
@@ -365,6 +365,13 @@ def require_tiled_view_type(var: Var) -> TiledViewTy:
     return ty
 
 
+def require_raw_array_memory_type(var: Var) -> RawArrayMemoryTy:
+    ty = var.get_type()
+    if not isinstance(ty, RawArrayMemoryTy):
+        raise _make_type_error(f"Expected a RawArrayMemory, but given value has type {ty}", var)
+    return ty
+
+
 def require_list_type(var: Var) -> ListTy:
     ty = var.get_type()
     if not isinstance(ty, ListTy):

src/cuda/tile/_ir/ops.py

@@ -25,7 +25,7 @@
     enter_nested_block, nested_block, PhiState, LoopVarState,
     TupleValue, make_aggregate, RangeValue, BoundMethodValue, ArrayValue, ConstantState,
     ListValue, TiledViewValue, ClosureValue, MemoryEffect, attribute, operand,
-    BlockRestriction, FormattedStringValue,
+    BlockRestriction, FormattedStringValue, RawArrayMemoryValue
 )
 from .type import PointerTy
 from . import hir
@@ -42,7 +42,7 @@
     require_optional_constant_str, PrintfValidator, require_tile_maybe_loose_type,
     require_0d_tile_maybe_loose_type, require_bool, require_optional_range_type,
     require_tile_or_tile_tuple_type, require_constant_scalar_tuple, require_constant_scalar,
-    require_callable_type)
+    require_callable_type, require_raw_array_memory_type)
 from .ops_utils import (
     BINOP_REGISTRY, UNARYOP_REGISTRY,
     check_rd_and_ftz, PaddingMode, get_default_order,
@@ -60,7 +60,7 @@
     ListTy, make_tile_ty, SliceType, DTypeConstructor, RangeIterType, Type,
     NONE, ModuleTy, TypeTy, LooselyTypedScalar, DTypeSpec, StringTy, InvalidType,
     array_size_type, ClosureTy, LiveCapturedScope, TokenTy, TiledViewTy, FormattedStringTy,
-    StringFormat, FormattedPiece,
+    StringFormat, FormattedPiece, RawArrayMemoryTy
 )
 from cuda.tile._datatype import (
     DType, is_integral, is_float, is_signed, is_boolean,
@@ -1593,6 +1593,7 @@ def getattr_impl(object: Var, name: Var) -> Var:
         case ArrayTy(), "strides": return build_tuple(object.get_aggregate().strides)
         case ArrayTy(), "slice": return bind_method(object, ct._m_array_slice)
         case ArrayTy(), "tiled_view": return bind_method(object, ct._m_array_tiled_view)
+        case ArrayTy(), "get_raw_memory": return bind_method(object, ct._m_array_get_raw_memory)
 
         case TileTy(), "dtype": return loosely_typed_const(ty.dtype)
         case TileTy(), "shape": return loosely_typed_const(ty.shape)
@@ -1614,6 +1615,12 @@ def getattr_impl(object: Var, name: Var) -> Var:
         case TiledViewTy(), "load": return bind_method(object, ct._m_tiled_view_load)
         case TiledViewTy(), "store": return bind_method(object, ct._m_tiled_view_store)
 
+        case RawArrayMemoryTy(), "dtype": return loosely_typed_const(ty.dtype)
+        case RawArrayMemoryTy(), "load_offset": return bind_method(
+            object, ct._m_raw_array_memory_load_offset)
+        case RawArrayMemoryTy(), "store_offset": return bind_method(
+            object, ct._m_raw_array_memory_store_offset)
+
         case ModuleTy(), _:
             try:
                 return loosely_typed_const(getattr(ty.py_mod, attr_name))
@@ -2154,6 +2161,83 @@ def _tile_load_impl_inner(array: Var, index_items: tuple[Var, ...], shape: Seque
     return reshape(result, shape)
 
 
+@impl(ct._m_array_get_raw_memory)
+def get_raw_memory_impl(array: Var) -> Var:
+    array_ty = require_array_type(array)
+    array_val = array.get_aggregate()
+    assert isinstance(array_val, ArrayValue)
+    base_ptr = array_val.base_ptr
+    raw_mem_ty = RawArrayMemoryTy(array_ty.dtype)
+    [ret] = unflatten_aggregates((base_ptr,), (raw_mem_ty,), (raw_mem_ty,))
+    return ret
+
+
+@impl(ct._m_raw_array_memory_load_offset)
+def raw_array_memory_load_offset_impl(raw_array_memory: Var, offset: Var, mask: Var,
+                                      padding_value: Var, latency: Var) -> Var:
+    raw_mem_ty = require_raw_array_memory_type(raw_array_memory)
+    raw_mem_val = raw_array_memory.get_aggregate()
+    assert isinstance(raw_mem_val, RawArrayMemoryValue)
+    base_ptr = raw_mem_val.base_ptr
+
+    offset = astype(offset, datatype.uint64)
+    pointer = pointer_offset(base_ptr, offset)
+    pointer_ty = pointer.get_type()
+    pointer_shape = pointer_ty.shape
+    array_dtype = raw_mem_ty.dtype
+
+    final_mask = _process_custom_mask(mask, None, pointer_shape)
+
+    if padding_value.is_constant() and padding_value.get_constant() is None:
+        padding_var: Optional[Var] = None
+    else:
+        padding_ty = require_tile_type(padding_value)
+        padding_shape = padding_ty.shape
+        if not is_shape_broadcastable_to(padding_shape, pointer_shape):
+            raise TileTypeError(f"Padding shape {padding_shape} is not broadcastable to the"
+                                f" offset shape {pointer_shape}")
+        padding_var = _implicit_cast(padding_value, array_dtype, "Invalid padding value")
+        padding_var = broadcast_to(padding_var, pointer_shape)
+
+    latency_val = require_optional_constant_int(latency)
+    _check_load_store_hints(latency_val)
+    result, _token = load_pointer(pointer, final_mask, padding_var, latency_val)
+    return result
+
+
+@impl(ct._m_raw_array_memory_store_offset)
+def raw_array_memory_store_offset_impl(raw_array_memory: Var, offset: Var, value: Var,
+                                       mask: Var, latency: Var) -> None:
+    raw_mem_ty = require_raw_array_memory_type(raw_array_memory)
+    raw_mem_val = raw_array_memory.get_aggregate()
+    assert isinstance(raw_mem_val, RawArrayMemoryValue)
+    base_ptr = raw_mem_val.base_ptr
+
+    offset = astype(offset, datatype.uint64)
+    pointer = pointer_offset(base_ptr, offset)
+    pointer_ty = pointer.get_type()
+    pointer_shape = pointer_ty.shape
+    array_dtype = raw_mem_ty.dtype
+
+    final_mask = _process_custom_mask(mask, None, pointer_shape)
+    value = _get_scatter_value(value, pointer_shape, array_dtype, "Value",
+                               array_name="RawArrayMemory")
+
+    latency_val = require_optional_constant_int(latency)
+    _check_load_store_hints(latency_val)
+    [_token] = store_pointer(pointer, value, final_mask, latency_val)
+
+
+def tile_load(array: Var, index: tuple[Var, ...], shape: Sequence[int], order: Sequence[int],
+              padding_mode: PaddingMode, latency: Optional[int],
+              allow_tma: Optional[bool]) -> tuple[Var, Var]:
+    res_ty = make_tile_ty(array.get_type().dtype, shape)
+    return add_operation(TileLoad, (res_ty, TokenTy()),
+                         array=array, index=index, order=tuple(order),
+                         padding_mode=padding_mode, latency=latency,
+                         allow_tma=allow_tma)
+
+
 @impl(ct.load)
 def tile_load_impl(array: Var, index: Var, shape: Var, order: Var,
                    padding_mode: Var, latency: Var, allow_tma: Var) -> Var:
@@ -2369,7 +2453,8 @@ def scatter_impl(array: Var, indices: Var, value: Var, mask: Var,
 
 
 def _get_scatter_value(value: Var, pointer_shape: Tuple[int, ...], array_dtype: DType,
-                       value_name: str, cast_dtype: bool = True) -> Var:
+                       value_name: str, cast_dtype: bool = True,
+                       array_name: str = "array") -> Var:
     value_ty = require_tile_type(value)
     value_shape = value_ty.shape
 
@@ -2379,7 +2464,7 @@ def _get_scatter_value(value: Var, pointer_shape: Tuple[int, ...], array_dtype:
 
     if cast_dtype:
         value = _implicit_cast(value, array_dtype,
-                               "Stored value is incompatible with array's dtype")
+                               f"Stored value is incompatible with {array_name}'s dtype")
     return broadcast_to(value, pointer_shape)
 
 

src/cuda/tile/_ir/type.py

@@ -465,6 +465,31 @@ def __str__(self):
                 f"padding_mode={self.padding_mode}]")
 
 
+# ============== Raw Array Memory Type ===============
+
+
+@dataclass(frozen=True)
+class RawArrayMemoryTy(Type):
+    """Type for a RawArrayMemory object that allows load/store by element offset (no index math)."""
+    dtype: "DType"
+
+    def is_aggregate(self) -> bool:
+        return True
+
+    def aggregate_item_types(self) -> tuple["Type", ...]:
+        base_ptr_ty = PointerTy(self.dtype)
+        base_ptr_tile_ty = TileTy(base_ptr_ty, TupleTy(()))
+        return (base_ptr_tile_ty,)
+
+    def make_aggregate_value(self, items: tuple["Var", ...]) -> "AggregateValue":
+        from .ir import RawArrayMemoryValue
+        assert len(items) == 1
+        return RawArrayMemoryValue(items[0])
+
+    def __str__(self):
+        return f"RawArrayMemory[{self.dtype}]"
+
+
 # ============== List Type ===============
 
 

src/cuda/tile/_stub.py

@@ -232,6 +232,64 @@ def tiled_view(self, tile_shape: Constant[Shape], *,
         """
         return _m_array_tiled_view(self, tile_shape, padding_mode=padding_mode)
 
+    def get_raw_memory(self) -> "RawArrayMemory":
+        """Returns an object that allows loading and storing by element offset.
+
+        The returned object holds the array's base pointer. Use
+        :py:meth:`RawArrayMemory.load_offset`
+        and :py:meth:`RawArrayMemory.store_offset` with an offset in **elements** (no shape/stride
+        index calculation). Useful when you already have memory offsets.
+
+        Returns:
+            RawArrayMemory:
+        """
+        return _m_array_get_raw_memory(self)
+
+
+class RawArrayMemory:
+    """Type stub for RawArrayMemory objects returned by :py:meth:`Array.get_raw_memory`."""
+
+    @property
+    @function
+    def dtype(self) -> "DType":
+        """The data type of the elements in the |RawArrayMemory|.
+
+        Returns:
+            DType (constant):
+        """
+
+    def load_offset(self, offset: "TileOrScalar", /, *,
+                    mask: Optional["Tile"] = None,
+                    padding_value: "TileOrScalar" = 0,
+                    latency: Optional[int] = None) -> "Tile":
+        """Loads from memory at base_ptr + offset (offset in elements).
+
+        Args:
+            offset: Element offset(s); scalar or tile of integer type.
+            mask: Optional boolean mask; where False, padding_value is used instead of load.
+            padding_value: Value used when mask is False; default 0.
+            latency: Optional latency hint (1--10).
+
+        Returns:
+            Tile: Loaded tile; shape matches broadcast(offset).
+        """
+        return _m_raw_array_memory_load_offset(
+            self, offset, mask=mask, padding_value=padding_value, latency=latency)
+
+    def store_offset(self, offset: "TileOrScalar", value: "TileOrScalar", /, *,
+                     mask: Optional["Tile"] = None,
+                     latency: Optional[int] = None) -> None:
+        """Stores to memory at base_ptr + offset (offset in elements).
+
+        Args:
+            offset: Element offset(s); scalar or tile of integer type.
+            value: Value(s) to store; broadcast to offset shape.
+            mask: Optional boolean mask; where False, no store occurs.
+            latency: Optional latency hint (1--10).
+        """
+        return _m_raw_array_memory_store_offset(
+            self, offset, value, mask=mask, latency=latency)
+
 
 class Tile:
     """Type stub for a |tile|."""
@@ -2430,3 +2488,22 @@ def _m_tiled_view_load(tiled_view, index, *, latency, allow_tma): ...
 @function
 def _m_tiled_view_store(tiled_view, index, tile, *, latency, allow_tma): ...
 # TiledView.store(index, tile, latency=latency, allow_tma=allow_tma)
+
+
+@function
+def _m_array_get_raw_memory(array: Array) -> RawArrayMemory: ...  # Array.get_raw_memory()
+
+
+@function
+def _m_raw_array_memory_load_offset(
+        raw_array_memory: RawArrayMemory, offset: TileOrScalar, /, *,
+        mask: Optional[Tile] = None,
+        padding_value: TileOrScalar = 0,
+        latency: Optional[int] = None) -> Tile: ...  # RawArrayMemory.load_offset()
+
+
+@function
+def _m_raw_array_memory_store_offset(
+        raw_array_memory: RawArrayMemory, offset: TileOrScalar, value: TileOrScalar, /, *,
+        mask: Optional[Tile] = None,
+        latency: Optional[int] = None) -> None: ...  # RawArrayMemory.store_offset()