[lang] Support vector astype

Signed-off-by: Qiqi Xiao <qiqix@nvidia.com>

Author: xiaoqiqi177

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

@@ -8,6 +8,7 @@
     ImplRegistry,
     require_dtype_spec,
 )
+from cuda.tile._ir.arithmetic_ops import astype
 from cuda.tile._ir.cast_ops import implicit_cast
 from cuda.tile._ir.core_ops import bind_method, build_tuple, loosely_typed_const
 from cuda.tile._ir.ops import strictly_typed_const
@@ -175,6 +176,16 @@ def vector_with_item_impl(
     return vector_with_item(self, index, value)
 
 
+@impl(getattr, overload=(VectorTy, "astype"))
+def getattr_vector_astype(object: Var[VectorTy], name: Var):
+    return bind_method(object, Vector.astype)
+
+
+@impl(Vector.astype)
+def vector_astype_impl(self: Var[VectorTy], dtype: Var) -> Var[VectorTy]:
+    return astype(self, require_dtype_spec(dtype))
+
+
 @impl(operator.getitem, overload=(VectorTy, WILDCARD))
 def vector_getitem(object: Var[VectorTy], key: Var[ScalarTy]) -> Var[ScalarTy]:
     result_dtype = object.get_type().element_dtype

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

@@ -143,6 +143,16 @@ def with_item(self, index: int, value: T) -> "Vector[T]":
             value: New value.
         """
 
+    @stub
+    def astype(self, dtype: "DType") -> "Vector":
+        """Convert each element to ``dtype``.
+
+        Returns a new vector of the same length with the given dtype.
+
+        Args:
+            dtype: Target data type of the result vector.
+        """
+
 
 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

@@ -73,6 +73,20 @@ def kernel(out):
     assert out.cpu().tolist() == [1, 1, 1, 4]
 
 
+def test_astype_on_vector():
+    @cl.kernel
+    def kernel(inp, out):
+        vector = inp.get_base_pointer().load(count=4, alignment=16)
+        halved = vector.astype(cl.float16)
+        out.get_base_pointer().store(halved, alignment=8)
+
+    values = [1.0, 2.0, 3.0, 4.0]
+    inp = torch.tensor(values, dtype=torch.float32).cuda()
+    out = torch.zeros(4, dtype=torch.float16).cuda()
+    cl.launch(torch.cuda.current_stream(), (1,), (1,), kernel, (inp, out))
+    assert out.cpu().tolist() == values
+
+
 @pytest.mark.parametrize('length', (2, 4))
 def test_vector_tuple(length):
     expect = tuple(range(length))

experimental/cuda-lang/tutorial/fp16_gemm_0.py

@@ -89,11 +89,11 @@ def _to_float16_vector(values, base, vsize):
     """Convert one FP32 vector slice to FP16."""
     return cl.Vector(
         *tuple(
-            cl.float16(values[base + i])
+            values[base + i]
             for i in cl.static_iter(range(vsize))
         ),
-        dtype=cl.float16,
-    )
+        dtype=cl.float32,
+    ).astype(cl.float16)
 
 
 @cl.kernel

experimental/cuda-lang/tutorial/fp16_gemm_1.py

@@ -64,11 +64,11 @@ def _to_float16_vector(values, base, vsize):
     """Convert one FP32 vector slice to FP16."""
     return cl.Vector(
         *tuple(
-            cl.float16(values[base + i])
+            values[base + i]
             for i in cl.static_iter(range(vsize))
         ),
-        dtype=cl.float16,
-    )
+        dtype=cl.float32,
+    ).astype(cl.float16)
 
 
 @cl.kernel

experimental/cuda-lang/tutorial/fp16_gemm_3.py

@@ -80,11 +80,11 @@ def _to_float16_vector(values, base, vsize):
     """Convert one FP32 vector slice to FP16."""
     return cl.Vector(
         *tuple(
-            cl.float16(values[base + i])
+            values[base + i]
             for i in cl.static_iter(range(vsize))
         ),
-        dtype=cl.float16,
-    )
+        dtype=cl.float32,
+    ).astype(cl.float16)
 
 
 @dataclass(frozen=True)

experimental/cuda-lang/tutorial/nvfp4_gemm_1_ws.py

@@ -81,11 +81,11 @@ def _to_float16_vector(
 ) -> cl.Vector[cl.float16]:
     return cl.Vector(
         *tuple(
-            cl.float16(values[base + i])
+            values[base + i]
             for i in cl.static_iter(range(count))
         ),
-        dtype=cl.float16,
-    )
+        dtype=cl.float32,
+    ).astype(cl.float16)
 
 
 @cl.kernel

experimental/cuda-lang/tutorial/nvfp4_gemm_2_quantize_fp4.py

@@ -138,9 +138,9 @@ def _as_float32_vector_64(regs):
 
 def _to_float16_vector(values, base, count):
     return cl.Vector(
-        *tuple(cl.float16(values[base + i]) for i in cl.static_iter(range(count))),
-        dtype=cl.float16,
-    )
+        *tuple(values[base + i] for i in cl.static_iter(range(count))),
+        dtype=cl.float32,
+    ).astype(cl.float16)
 
 
 def _pack_fp32x8_to_e2m1x8(values, base):