[lang] Introduce cl.Array.from_parts

Replaces `cl.reinterpret_pointer_as_array` with `Array.from_parts`.

ai-use: updating callsites.

Signed-off-by: Asher Mancinelli <amancinelli@nvidia.com>

Author: ashermancinelli

experimental/cuda-lang/src/cuda/lang/__init__.py

@@ -155,7 +155,6 @@
     map_shared_to_cluster,
     map_shared_to_leader_block,
     shared_cluster_leader_bit_mask,
-    reinterpret_pointer_as_array,
     _inline_ptx,
     ptx_comment,
     shfl_sync,
@@ -361,7 +360,6 @@
     "map_shared_to_cluster",
     "map_shared_to_leader_block",
     "shared_cluster_leader_bit_mask",
-    "reinterpret_pointer_as_array",
     "_inline_ptx",
     "ptx_comment",
     "Array",

experimental/cuda-lang/src/cuda/lang/_ir/op_defs.py

@@ -205,11 +205,6 @@ def has_observable_effect(self) -> bool:
         return self.mmio or self.memory_order is MemoryOrder.ACQUIRE
 
 
-@dataclass(eq=False)
-class ReinterpretPointerAsArray(Operation, opcode="reinterpret_ptr_as_array"):
-    pointer: Var = operand()
-
-
 @dataclass
 class TensorMapAsOpaquePtr(Operation, opcode="tensor_map_as_opaque_ptr"):
     tensor_map: Var = operand()

experimental/cuda-lang/src/cuda/lang/_ir/op_impl/pointer_impl.py

@@ -25,7 +25,6 @@
     AtomicLoad,
     AtomicStore,
     LoadPointer,
-    ReinterpretPointerAsArray,
     StorePointer,
 )
 from cuda.lang._ir.type_checking_helpers import (
@@ -44,9 +43,9 @@
     uint64,
 )
 from cuda.tile._ir.arithmetic_ops import astype, binary_arithmetic_tensorlike_raw
-from cuda.tile._ir.cast_ops import address_space_cast, implicit_cast, reinterpret_pointer
+from cuda.tile._ir.cast_ops import address_space_cast, implicit_cast
 from cuda.tile._ir.core_ops import bind_method, loosely_typed_const, strictly_typed_const
-from cuda.tile._ir.ir import add_operation_variadic
+from cuda.tile._ir.ir import add_operation_variadic, make_aggregate
 from cuda.tile._ir.op_impl import (
     ImplRegistry,
     WILDCARD,
@@ -460,39 +459,29 @@ def address_space_cast_impl(value: Var, memory_space: Var) -> Var:
     return address_space_cast(value, memory_space)
 
 
-@impl(core_api.reinterpret_pointer_as_array)
-def reinterpret_pointer_as_array_impl(pointer: Var, dtype: Var, shape: Var, strides: Var) -> Var:
-    if not strides.is_constant() or strides.get_constant() is not None:
-        raise TypeCheckingError(
-            "Reinterpreting a pointer as an array with "
-            "non-default strides is not yet implemented."
-        )
-    pointer_ty = require_pointer_type(pointer)
+@impl(Array.from_parts)
+def array_from_parts_impl(pointer: Var, shape: Var, strides: Var) -> Var:
+    pointer_ty = require_concrete_pointer_type(pointer)
     shape = require_constant_int_tuple(shape, allow_single_int=True)
-    dtype = require_dtype_spec(dtype)
-    strides = contiguous_strides_from_shape(shape)
-    memory_space = pointer_ty.memory_space
-
-    typed_pointer_ty = PointerTy(pointer_dtype(dtype, memory_space))
-    if pointer.get_type() != typed_pointer_ty:
-        pointer = reinterpret_pointer(pointer, typed_pointer_ty.pointer_dtype)
+    if strides.is_constant() and strides.get_constant() is None:
+        strides = contiguous_strides_from_shape(shape)
+    else:
+        strides = require_constant_int_tuple(strides, allow_single_int=True)
+        if len(strides) != len(shape):
+            raise TypeCheckingError(
+                f"Shape and strides must have the same rank, got {len(shape)} and {len(strides)}"
+            )
     index_dtype = datatype.int32
     array_ty = ArrayTy(
-        dtype,
+        pointer_ty.pointee_dtype,
         shape=shape,
         strides=strides,
         typing_hooks=pointer.ctx.typing_hooks,
         index_dtype=index_dtype,
-        memory_space=memory_space,
-    )
-    result = add_operation(
-        ReinterpretPointerAsArray,
-        array_ty,
-        pointer=pointer,
+        memory_space=pointer_ty.memory_space,
     )
-    # FIXME: it seems that the index dtype should be derived from the dtype of shape/strides instead
+    # FIXME: Derive the index dtype from the shape and stride dtypes.
     size_ty = ScalarTy(index_dtype)
     shape_vars = tuple(strictly_typed_const(extent, size_ty) for extent in shape)
     stride_vars = tuple(strictly_typed_const(extent, size_ty) for extent in strides)
-    result.set_aggregate(ArrayValue(pointer, shape_vars, stride_vars))
-    return result
+    return make_aggregate(ArrayValue(pointer, shape_vars, stride_vars), array_ty)

experimental/cuda-lang/src/cuda/lang/_ir/ops.py

@@ -105,7 +105,6 @@
     VectorReduce,
     StorePointer,
     LoadPointer,
-    ReinterpretPointerAsArray,
     BitCast,
     FmaOperation,
 )
@@ -958,7 +957,6 @@ def _unflatten_aggregate_array_impl(val: ArrayValue, ty: ArrayTy, result_var: Va
     "PointerOffset",
     "LoadPointer",
     "ReinterpretPointer",
-    "ReinterpretPointerAsArray",
     "StorePointer",
     "RawWhereOperation",
     "Unary",

experimental/cuda-lang/src/cuda/lang/_passes/handle_dyn_shared_mem.py

@@ -7,7 +7,8 @@
 from cuda.lang._ir._host_program import HostProgram, host_program_to_ir
 from cuda.lang._ir.ops import AllocDynSharedMemory, GetDynSharedMemoryBasePtr, \
     get_dyn_shared_memory_base_ptr
-from cuda.lang._ir.op_impl.pointer_impl import reinterpret_pointer, pointer_with_offset
+from cuda.lang._ir.op_impl.pointer_impl import pointer_with_offset
+from cuda.tile._ir.cast_ops import reinterpret_pointer
 from cuda.lang._datatype import int32
 from cuda.lang._exception import TypeCheckingError
 from cuda.lang._ir.type import PointerTy, ScalarTy

experimental/cuda-lang/src/cuda/lang/_passes/ir2mlir/pass_definition.py

@@ -353,13 +353,6 @@ def _get_mlir_comparison_op(
                 return None
 
 
-# These operations have aggregate results. The RHS's elements are stored to
-# the LHS's when lowering Assign operations and are no-ops at the MLIR level.
-_NOOP_LOWERINGS = frozenset([
-    ops.ReinterpretPointerAsArray,
-])
-
-
 @dataclass(kw_only=True)
 class MLIRLoweringContext:
     """Mutable data shared by host and device MLIR operation lowering."""
@@ -945,9 +938,6 @@ def lower_block(self, ir_block: ir.Block) -> None:
             for operation in ir_block.operations:
                 context.current_op = operation
 
-                if isinstance(operation, tuple(_NOOP_LOWERINGS)):
-                    continue
-
                 operation_loc = None
                 if self.compiler_options.debug_info == "line":
                     operation_loc = ir_loc_to_mlir_location(operation.loc)

experimental/cuda-lang/src/cuda/lang/_stub/core_api.py

@@ -34,6 +34,47 @@ class Array(TileArray, Generic[T]):
     N-dimensional array type.
     """
 
+    @staticmethod
+    @stub
+    def from_parts(
+        pointer: Pointer[T],
+        shape: int | tuple[int, ...],
+        strides: int | tuple[int, ...] | None = None,
+    ) -> "Array[T]":
+        """Create an array from a pointer, shape, and optional strides.
+
+        The pointer sets the array data type, so the pointer must be typed.
+        If ``strides`` is ``None``, the array uses contiguous row-major strides.
+
+        Args:
+            pointer: Typed pointer to the first array element.
+            shape: Number of elements in each array dimension.
+            strides: Step in elements for each dimension.
+
+        Returns:
+            Array that refers to the memory at ``pointer``.
+
+        Examples:
+
+        .. testcode::
+            :template: kernel_wrapper.py
+
+            smem_array = cl.shared_array(1, cl.int32)
+            smem_array[0] = 5
+
+            smem_ptr = smem_array.pointer()
+            smem_array_2 = cl.Array.from_parts(smem_ptr, 1)
+
+            # Assignment through the new array changes memory referenced by the
+            # original array.
+            smem_array_2[0] = 7
+            print(smem_array[0])
+
+        .. testoutput::
+
+            7
+        """
+
     @stub
     @property
     def dtype(self): ...
@@ -1077,46 +1118,6 @@ def map_shared_to_leader_block(pointer: Pointer[T]) -> Pointer[T]:
     """
 
 
-@stub
-def reinterpret_pointer_as_array(
-    pointer: Pointer[T],
-    dtype: DType,
-    shape: tuple[int, ...],
-    strides: tuple[int, ...] | None = None,
-) -> Array[T]:
-    """
-    Args:
-        pointer: Pointer[T]
-        dtype: DType
-        shape: tuple[int, ...]
-        strides: tuple[int, ...] | None = None
-
-    Returns:
-        Array with the specified base pointer, dtype, shape, and strides.
-
-    Examples:
-
-    .. testcode::
-        :template: kernel_wrapper.py
-
-        smem_array = cl.shared_array(1, cl.int32)
-        smem_array[0] = 5
-
-        smem_ptr = smem_array.pointer()
-        smem_array_2 = cl.reinterpret_pointer_as_array(smem_ptr, shape=1, dtype=cl.int32)
-
-        # Assignment through the reconstructed array is equivalent
-        # to assignment through the original.
-        smem_array_2[0] = 7
-        print(smem_array[0])
-
-    .. testoutput::
-
-        7
-
-    """
-
-
 def nanosleep(nanoseconds: int):
     """
     Sleep for ``nanoseconds`` nanoseconds.

experimental/cuda-lang/test/examples/test_histogram64.py

@@ -55,10 +55,8 @@ def histogram64_kernel(d_PartialHistograms, d_Data, dataCount):
         shape=(HISTOGRAM64_THREADBLOCK_SIZE * HISTOGRAM64_BIN_COUNT,),
         dtype=cl.uint8,
     )
-    s_ThreadBase = cl.reinterpret_pointer_as_array(
-        s_Hist.pointer((threadPos,)),
-        cl.uint8,
-        1,
+    s_ThreadBase = cl.Array.from_parts(
+        s_Hist.pointer(threadPos), 1
     )
 
     for i in cl.static_iter(range(HISTOGRAM64_BIN_COUNT)):
@@ -76,10 +74,8 @@ def histogram64_kernel(d_PartialHistograms, d_Data, dataCount):
     cl.barrier_sync_block()
 
     if tx < HISTOGRAM64_BIN_COUNT:
-        s_HistBase = cl.reinterpret_pointer_as_array(
-            s_Hist.pointer((tx * HISTOGRAM64_THREADBLOCK_SIZE,)),
-            cl.uint8,
-            1,
+        s_HistBase = cl.Array.from_parts(
+            s_Hist.pointer(tx * HISTOGRAM64_THREADBLOCK_SIZE), 1
         )
         sum = cl.uint32(0)
         pos = 4 * (tx & (SHARED_MEMORY_BANKS - 1))

experimental/cuda-lang/test/test_nvvm_intrinsics.py

@@ -76,7 +76,8 @@ def test_smem_pointer_arg_intrinsic():
     def kern(y):
         smem = cl.shared_array(1, dtype=cl.int32)
         p = cl._nvvm.mapa_shared_cluster(smem.pointer(), 0)
-        a = cl.reinterpret_pointer_as_array(p, cl.int32, (1,))
+        p = cl.bitcast(p, cl.pointer_dtype(cl.int32, p.memory_space))
+        a = cl.Array.from_parts(p, 1)
         a[0] = 13
         y[0] = smem[0]
 
@@ -91,7 +92,8 @@ def test_generic_pointer_arg_intrinsic():
     def kern(y):
         smem = cl.shared_array(1, dtype=cl.int32)
         p = cl._nvvm.mapa(smem.pointer(), 0)
-        a = cl.reinterpret_pointer_as_array(p, cl.int32, (1,))
+        p = cl.bitcast(p, cl.pointer_dtype(cl.int32, p.memory_space))
+        a = cl.Array.from_parts(p, 1)
         a[0] = 13
         y[0] = smem[0]
 

experimental/cuda-lang/test/test_pointers.py

@@ -382,7 +382,7 @@ def test_ptr_roundtrip():
     def kernel(A):
         B = cl.shared_array(shape=(3, 3), dtype=cl.int32)
         smem = B.pointer()
-        B2 = cl.reinterpret_pointer_as_array(smem, cl.int32, 1)
+        B2 = cl.Array.from_parts(smem, 1)
         B2[0] = 1
         A[0] = B[0, 0]
         B2[0] = 2
@@ -393,6 +393,68 @@ def kernel(A):
     assert A.cpu().tolist() == [1, 2]
 
 
+def test_array_from_parts_with_strides():
+    @cl.kernel
+    def kernel(array):
+        view = cl.Array.from_parts(array.pointer(), (2, 2), (3, 1))
+        cl.static_assert(view.shape == (2, 2))
+        cl.static_assert(view.strides == (3, 1))
+        cl.static_assert(view.dtype == cl.int32)
+        view[0, 0] = 10
+        view[0, 1] = 11
+        view[1, 0] = 12
+        view[1, 1] = 13
+
+    array = torch.zeros(5, dtype=torch.int32, device="cuda")
+    cl.launch(torch.cuda.current_stream(), (1,), (1,), kernel, (array,))
+    assert array.cpu().tolist() == [10, 11, 0, 12, 13]
+
+
+def test_array_from_parts_uses_pointer_dtype():
+    @cl.kernel
+    def kernel(array):
+        pointer = cl.bitcast(
+            array.pointer(),
+            cl.pointer_dtype(cl.float32, cl.MemorySpace.GLOBAL),
+        )
+        view = cl.Array.from_parts(pointer, 1)
+        cl.static_assert(view.dtype == cl.float32)
+        view[0] = 1.5
+
+    array = torch.zeros(1, dtype=torch.int32, device="cuda")
+    cl.launch(torch.cuda.current_stream(), (1,), (1,), kernel, (array,))
+    assert array.cpu().item() == 0x3FC00000
+
+
+def test_array_from_parts_rejects_stride_rank_mismatch():
+    def kernel(array):
+        cl.Array.from_parts(array.pointer(), (2, 2), (1,))
+
+    compile_kernel(
+        kernel,
+        signature=KernelSignature((make_symbolic_tensor(4, cl.int32),)),
+        raises=pytest.raises(
+            TypeCheckingError,
+            match="Shape and strides must have the same rank, got 2 and 1",
+        ),
+    )
+
+
+def test_array_from_parts_rejects_opaque_pointer():
+    def kernel(array):
+        pointer = cl.bitcast(array.pointer(), cl.opaque_pointer_dtype())
+        cl.Array.from_parts(pointer, 1)
+
+    compile_kernel(
+        kernel,
+        signature=KernelSignature((make_symbolic_tensor(1, cl.int32),)),
+        raises=pytest.raises(
+            TypeCheckingError,
+            match="Expected concrete pointer type but got opaque_pointer",
+        ),
+    )
+
+
 def test_pointer_smem():
     @cl.kernel
     def kernel(A):