From dfe7dea0ebb6412101192e60f2a0606d9ec24382 Mon Sep 17 00:00:00 2001 From: Samridh <109811970+samridh10exe@users.noreply.github.com> Date: Thu, 15 Oct 2026 00:41:08 -0400 Subject: [PATCH] [lang] Fix Array.slice with integer bounds Signed-off-by: Samridh <109811970+samridh10exe@users.noreply.github.com> --- .../src/cuda/lang/_ir/op_impl/pointer_impl.py | 27 +++++++++++ experimental/cuda-lang/test/test_slice.py | 46 +++++++++++++++++++ 2 files changed, 73 insertions(+) diff --git a/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/pointer_impl.py b/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/pointer_impl.py index 9fb8b6d9..0db8f834 100644 --- a/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/pointer_impl.py +++ b/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: diff --git a/experimental/cuda-lang/test/test_slice.py b/experimental/cuda-lang/test/test_slice.py index 1ca658ac..06a3f283 100644 --- a/experimental/cuda-lang/test/test_slice.py +++ b/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,))