[lang] Support replacing vector elements

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

Author: ashermancinelli

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

@@ -9,6 +9,7 @@
     require_dtype_spec,
 )
 from cuda.tile._ir.cast_ops import implicit_cast
+from cuda.tile._ir.core_ops import bind_method
 from cuda.tile._ir.ops import strictly_typed_const
 from cuda.tile._ir.ops_utils import promote_dtypes
 from cuda.tile._ir.type import LooselyTypedScalar
@@ -35,7 +36,7 @@ def vector_undef(res_type: Type):
     )
 
 
-def vector_setitem(vector: Var[VectorTy], key: int | Var[ScalarTy], value: Var[ScalarTy]):
+def vector_with_item(vector: Var[VectorTy], key: int | Var[ScalarTy], value: Var[ScalarTy]):
     ty = require_vector_type(vector)
     if isinstance(key, int):
         key = strictly_typed_const(key, ScalarTy(datatype.int32))
@@ -101,7 +102,7 @@ def vector_constructor_impl(elements: tuple[Var, ...], dtype: Var) -> Var[Vector
     res = vector_undef(VectorTy(element_dtype, len(elements)))
     for index, element in enumerate(elements):
         value = implicit_cast(element, element_dtype, f"Vector() element {index}")
-        res = vector_setitem(res, index, value)
+        res = vector_with_item(res, index, value)
     return res
 
 
@@ -130,20 +131,31 @@ def apply_one(i: int):
         )
 
     res = vector_undef(VectorTy(element_type.dtype, length))
-    res = vector_setitem(res, 0, first_element)
+    res = vector_with_item(res, 0, first_element)
     for i in range(1, length):
         element = apply_one(i)
-        res = vector_setitem(res, i, element)
+        res = vector_with_item(res, i, element)
 
     return res
 
 
-# the user can't call __setitem__ in kernel code because the semantics might be
-# confusing. We could expose an insertelement-like operation that maps to the
-# vector_setitem utility though.
 @impl(operator.setitem, overload=(VectorTy, WILDCARD, WILDCARD))
 def vector_setitem_impl(object: Var[VectorTy], key: Var, value: Var):
-    raise TypeCheckingError("Vectors are immutable: item assignment is not supported")
+    raise TypeCheckingError(
+        "Vectors are immutable. Consider calling vector.with_item() instead"
+    )
+
+
+@impl(getattr, overload=(VectorTy, "with_item"))
+def getattr_vector_with_item(object: Var[VectorTy], name: Var):
+    return bind_method(object, Vector.with_item)
+
+
+@impl(Vector.with_item)
+def vector_with_item_impl(
+    self: Var[VectorTy], index: Var[ScalarTy], value: Var[ScalarTy]
+) -> Var[VectorTy]:
+    return vector_with_item(self, index, value)
 
 
 @impl(operator.getitem, overload=(VectorTy, WILDCARD))

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

@@ -131,6 +131,18 @@ def __getitem__(self, item): ...
     @stub
     def __setitem__(self, key, value): ...
 
+    @stub
+    def with_item(self, index: int, value: T) -> "Vector[T]":
+        """Return a new vector with one element replaced.
+
+        Vectors have value semantics, so this operation does not modify the
+        original vector. ``index`` must select an element of the vector.
+
+        Args:
+            index: Index in vector to replace.
+            value: New value.
+        """
+
 
 class Pointer(Generic[T]):
     """Typed address into a CUDA memory space with low-level load and store operations."""

experimental/cuda-lang/test/test_vectors.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from test.util import compile_kernel
 import operator
 
 import pytest
@@ -399,16 +400,28 @@ def kernel(tensor):
 
 
 def test_vector_setitem():
-    @cl.kernel
     def kernel():
-        with cl.local_array(4, cl.int32) as arr:
-            v = arr.get_base_pointer().load(count=4)
-            v[0] = 1
+        v = cl.shared_array(1, cl.int8).get_base_pointer().load(count=2)
+        v[0] = 1
 
-    with pytest.raises(
-        TypeCheckingError, match="Vectors are immutable: item assignment is not supported"
-    ):
-        cl.compile_simt(kernel, [KernelSignature([])])
+    compile_kernel(
+        kernel,
+        raises=pytest.raises(TypeCheckingError, match="Vectors are immutable"),
+    )
+
+
+def test_vector_with_item():
+    @cl.kernel
+    def kernel(original, updated):
+        original_vector = original.get_base_pointer().load(count=4, alignment=16)
+        updated_vector = original_vector.with_item(2, 42)
+        updated.get_base_pointer().store(updated_vector, alignment=16)
+
+    a = torch.arange(4, dtype=torch.int32, device="cuda")
+    b = torch.arange(4, dtype=torch.int32, device="cuda")
+    cl.launch(torch.cuda.current_stream(), (1,), (1,), kernel, (a, b))
+    assert a.cpu().tolist() == [0, 1, 2, 3]
+    assert b.cpu().tolist() == [0, 1, 42, 3]
 
 
 def test_vector_from_tuple():