[lang] Fix Array.slice with integer bounds

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

@@ -58,9 +58,11 @@
 from cuda.tile._ir.op_impl import (
     ImplRegistry,
     WILDCARD,
+    normalize_axis,
     require_array_type,
     require_constant_bool,
     require_constant_enum,
+    require_constant_int,
     require_constant_pointer_info,
     require_constant_str,
     require_dtype_spec,
@@ -413,6 +415,31 @@ def array_getitem(object: Var, key: Var) -> Var:
     return array_from_parts_impl(pointer, build_tuple(new_shape), build_tuple(new_strides))
 
 
+@impl(Array.slice)
+def array_slice(self: Var, axis: Var, start: Var, stop: Var) -> Var:
+    array_ty = require_array_type(self)
+    array_val = self.get_aggregate()
+    axis = normalize_axis(require_constant_int(axis), array_ty.ndim)
+    start = implicit_cast(start, array_ty.index_dtype, "Invalid slice start")
+    stop = implicit_cast(stop, array_ty.index_dtype, "Invalid slice stop")
+    _require_nonnegative_constant(start, "Slice start must be non-negative")
+    _require_nonnegative_constant(stop, "Slice stop must be non-negative")
+    if start.is_constant() and stop.is_constant() and stop.get_constant() < start.get_constant():
+        raise TypeCheckingError("Slice stop must be greater than or equal to start")
+
+    stride = array_val.strides[axis]
+    offset = binary_arithmetic_tensorlike(
+        "mul", astype(start, datatype.uint64), astype(stride, datatype.uint64)
+    )
+    base_ptr = _get_array_base_pointer(self)
+    pointer = add_operation(PointerOffset, base_ptr.get_type(), pointer=base_ptr, offset=offset)
+    new_shape = tuple(
+        binary_arithmetic_tensorlike("sub", stop, start) if i == axis else dim
+        for i, dim in enumerate(array_val.shape)
+    )
+    return array_from_parts_impl(pointer, build_tuple(new_shape), build_tuple(array_val.strides))
+
+
 def load_pointer(pointer: Var[PointerTy], count: int | None = None, alignment: int | None = None):
     pointee_dtype = pointer.get_type().pointee_dtype
     if count is None or count == 1:

experimental/cuda-lang/test/test_slice.py

@@ -168,3 +168,49 @@ def kernel(a):
     a = torch.arange(16, dtype=torch.int32).reshape(4, 4).cuda()
     with pytest.raises(TypeCheckingError, match="The provided slice index is not valid"):
         cl.launch(torch.cuda.current_stream(), (1,), (1,), kernel, (a,))
+
+
+def test_array_slice_method():
+    @cl.kernel
+    def kernel(a, b):
+        view = a.slice(axis=0, start=4, stop=8)
+        cl.static_assert(view.shape[0] == 4)
+        b[0] = view[0]
+        b[1] = view[1]
+        b[2] = view[2]
+        b[3] = view[3]
+
+    a = torch.arange(16, dtype=torch.int32).cuda()
+    b = torch.zeros(4, dtype=torch.int32).cuda()
+    cl.launch(torch.cuda.current_stream(), (1,), (1,), kernel, (a, b))
+    assert b.cpu().tolist() == [4, 5, 6, 7]
+
+
+def test_array_slice_method_runtime_bounds():
+    @cl.kernel
+    def kernel(a, b, i):
+        view = a.slice(axis=-1, start=i[0], stop=i[1])
+        b[0] = view[0, 0]
+        b[1] = view[0, 1]
+        b[2] = view[1, 0]
+        b[3] = view.shape[1]
+
+    a = torch.arange(16, dtype=torch.int32).reshape((2, 8)).cuda()
+    b = torch.zeros(4, dtype=torch.int32).cuda()
+    i = torch.tensor([2, 5], dtype=torch.int32).cuda()
+    cl.launch(torch.cuda.current_stream(), (1,), (1,), kernel, (a, b, i))
+    assert b.cpu().tolist() == [2, 3, 10, 3]
+
+
+@pytest.mark.parametrize("start, stop, message",
+                         [(-1, 4, "Slice start must be non-negative"),
+                          (0, -1, "Slice stop must be non-negative"),
+                          (4, 2, "Slice stop must be greater than or equal to start")])
+def test_array_slice_method_static_reject(start, stop, message):
+    @cl.kernel
+    def kernel(a):
+        a.slice(axis=0, start=start, stop=stop)
+
+    a = torch.arange(16, dtype=torch.int32).cuda()
+    with pytest.raises(TypeCheckingError, match=message):
+        cl.launch(torch.cuda.current_stream(), (1,), (1,), kernel, (a,))