[lang] Fix cta group kind in cp_async_bulk_tensor hir2ir lowering

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

Author: ashermancinelli

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

@@ -19,7 +19,7 @@
     require_uniform_int_tuple_type,
     tensor_map_descriptor_like,
 )
-from cuda.tile._ir.op_impl import require_constant_enum
+from cuda.tile._ir.op_impl import require_constant_enum, require_optional_constant_enum
 import cuda.lang._mlir.nvvm as mlir
 
 
@@ -91,6 +91,9 @@ def cp_async_bulk_tensor_global_to_shared_impl(
         require_optional(multicast_mask, require_integral_scalar_type)
         require_optional(l2_cache_hint, require_integral_scalar_type)
         require_optional(predicate, require_boolean_scalar_type)
+        group_value = require_optional_constant_enum(group, cp_async.CTAGroup)
+        if group_value is not None:
+            group = loosely_typed_const(getattr(mlir.CTAGroupKind, group_value.name))
 
     return _raw_nvvm_mlir_operation_impl(
         nvvm_mlir_interfaces.cp_async_bulk_tensor_shared_cluster_global,

experimental/cuda-lang/test/test_cp_async.py

@@ -60,6 +60,61 @@ def kernel(x, pred, i, j, H: cl.Constant[int], W: cl.Constant[int]):
         )
         self.check_ptx_source(kernel, expect)
 
+    @require_blackwell_cc100()
+    @pytest.mark.parametrize(
+        "group,expect_group",
+        (
+            (cl.CTAGroup.CTA_1, "cta_group::1"),
+            (cl.CTAGroup.CTA_2, "cta_group::2"),
+        ),
+    )
+    def test_shared_cluster_group(self, group, expect_group):
+        @cl.kernel
+        def kernel(x, pred, i, j, H: cl.Constant[int], W: cl.Constant[int]):
+            tensor_map = cl.tensor_map_tiled(x, (H, W)).as_opaque_ptr()
+            smem = cl.shared_array(shape=(H * W,), dtype=cl.int32, alignment=512)
+            smem = cl.map_shared_to_cluster(smem.get_base_pointer(), 0)
+            mbar = cl.shared_array(1, cl.mbarrier, alignment=8).get_base_pointer()
+
+            cl.cp_async_bulk_tensor_global_to_shared(
+                tensor_map,
+                (i, j),
+                smem,
+                mbar,
+                group=group,
+            )
+
+        self.check_ptx_source(
+            kernel,
+            "cp.async.bulk.tensor.2d.shared::cluster.global",
+            expect_group,
+        )
+
+    @require_blackwell_cc100()
+    def test_shared_cluster_group_with_predicate_and_multicast(self):
+        @cl.kernel
+        def kernel(x, pred, i, j, H: cl.Constant[int], W: cl.Constant[int]):
+            tensor_map = cl.tensor_map_tiled(x, (H, W)).as_opaque_ptr()
+            smem = cl.shared_array(shape=(H * W,), dtype=cl.int32, alignment=512)
+            smem = cl.map_shared_to_cluster(smem.get_base_pointer(), 0)
+            mbar = cl.shared_array(1, cl.mbarrier, alignment=8).get_base_pointer()
+
+            cl.cp_async_bulk_tensor_global_to_shared(
+                tensor_map,
+                (i, j),
+                smem,
+                mbar,
+                multicast_mask=0x3,
+                group=cl.CTAGroup.CTA_2,
+                predicate=pred,
+            )
+
+        self.check_ptx_source(
+            kernel,
+            "cp.async.bulk.tensor.2d.shared::cluster.global",
+            "multicast::cluster",
+        )
+
     def test_unsupported_kwargs_for_cta_mode(self, subtests):
         @cl.kernel
         def k1(x, pred, i, j, H: cl.Constant[int], W: cl.Constant[int]):