[TEST] Cover Gluon Blackwell two-CTA tcgen05 ops

tests/end_to_end/test_gluon_blackwell_tcgen05_two_cta_ops.py

@@ -0,0 +1,125 @@
+import torch
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+from triton.experimental.gluon.language.nvidia import blackwell, hopper
+from triton.experimental.gluon.language.nvidia.hopper import mbarrier, tma
+from triton.experimental.gluon.nvidia.hopper import TensorDescriptor
+
+import triton_viz
+from triton_viz.core.callbacks import ForLoopCallbacks, OpCallbacks
+from triton_viz.core.client import Client
+
+
+class _NoOpClient(Client):
+    NAME = "noop"
+
+    def pre_run_callback(self, fn):
+        return True
+
+    def post_run_callback(self, fn):
+        return True
+
+    def arg_callback(self, name, arg, arg_cvt):
+        pass
+
+    def grid_callback(self, grid):
+        pass
+
+    def grid_idx_callback(self, grid_idx):
+        pass
+
+    def register_op_callback(self, op_type, *args, **kwargs):
+        return OpCallbacks()
+
+    def register_for_loop_callback(self):
+        return ForLoopCallbacks()
+
+    def finalize(self):
+        return []
+
+    def pre_warmup_callback(self, jit_fn, *args, **kwargs):
+        return True
+
+    def post_warmup_callback(self, jit_fn, ret):
+        pass
+
+
+@gluon.jit
+def _two_cta_tcgen05_kernel(
+    a_desc,
+    b_desc,
+    out_desc,
+    num_warps: gl.constexpr,
+):
+    gl.static_assert(gl.num_ctas() == 2)
+    cluster_m: gl.constexpr = a_desc.block_shape[0]
+    tile_n: gl.constexpr = b_desc.block_shape[1]
+    cta_m: gl.constexpr = cluster_m // 2
+    cga_layout: gl.constexpr = out_desc.layout.cga_layout
+
+    smem_a = gl.allocate_shared_memory(a_desc.dtype, a_desc.block_shape, a_desc.layout)
+    smem_b = gl.allocate_shared_memory(b_desc.dtype, b_desc.block_shape, b_desc.layout)
+
+    tma_barrier = gl.allocate_shared_memory(gl.int64, [1], mbarrier.MBarrierLayout())
+    mma_barrier = gl.allocate_shared_memory(gl.int64, [1], mbarrier.MBarrierLayout())
+    mbarrier.init(tma_barrier, count=1)
+    mbarrier.init(mma_barrier, count=1)
+
+    mbarrier.expect(tma_barrier, a_desc.nbytes_per_cta + b_desc.nbytes_per_cta)
+    tma.async_load(a_desc, [0, 0], tma_barrier, smem_a)
+    tma.async_load(b_desc, [0, 0], tma_barrier, smem_b)
+    mbarrier.wait(tma_barrier, phase=0, deps=[smem_a, smem_b])
+    mbarrier.invalidate(tma_barrier)
+
+    acc_layout: gl.constexpr = blackwell.TensorMemoryLayout(
+        block=(cta_m, tile_n),
+        col_stride=1,
+        cga_layout=cga_layout,
+        two_ctas=True,
+    )
+    acc = blackwell.allocate_tensor_memory(gl.float32, [cluster_m, tile_n], acc_layout)
+    blackwell.tcgen05_mma(smem_a, smem_b, acc, use_acc=False, mbarriers=[mma_barrier])
+    mbarrier.wait(mma_barrier, phase=0, deps=[smem_a, smem_b])
+    mbarrier.invalidate(mma_barrier)
+
+    out_smem = gl.allocate_shared_memory(
+        out_desc.dtype,
+        out_desc.block_shape,
+        out_desc.layout,
+    )
+    acc_reg_layout: gl.constexpr = acc.get_reg_layout(num_warps=num_warps)
+    out_smem.store(acc.load(acc_reg_layout).to(out_desc.dtype))
+    hopper.fence_async_shared()
+    tma.async_store(out_desc, [0, 0], out_smem)
+    tma.store_wait(0)
+
+
+def test_gluon_blackwell_two_cta_tcgen05_mma_on_cpu():
+    torch.manual_seed(12)
+    a = torch.randn(128, 16, dtype=torch.float16) / 4
+    b = torch.randn(16, 64, dtype=torch.float16) / 4
+    out = torch.empty((128, 64), dtype=torch.float16)
+    a_layout = gl.NVMMASharedLayout.get_default_for(
+        list(a.shape),
+        gl.float16,
+        cga_layout=[(1, 0)],
+    )
+    b_layout = gl.NVMMASharedLayout.get_default_for(
+        list(b.shape),
+        gl.float16,
+        cga_layout=[(0, 1)],
+    )
+    out_layout = gl.NVMMASharedLayout.get_default_for(
+        list(out.shape),
+        gl.float16,
+        cga_layout=[(1, 0)],
+    )
+    a_desc = TensorDescriptor.from_tensor(a, list(a.shape), a_layout)
+    b_desc = TensorDescriptor.from_tensor(b, list(b.shape), b_layout)
+    out_desc = TensorDescriptor.from_tensor(out, list(out.shape), out_layout)
+    kernel = triton_viz.trace(_NoOpClient(), frontend="gluon")(_two_cta_tcgen05_kernel)
+
+    kernel[(1,)](a_desc, b_desc, out_desc, num_warps=4, num_ctas=2)
+
+    expected = (a.float() @ b.float()).half()
+    torch.testing.assert_close(out, expected, atol=1e-3, rtol=1e-3)

triton_viz/core/simulation/gluon.py

@@ -216,7 +216,7 @@ def num_warps(self, _generator: Any = None):
         return 1
 
     def num_ctas(self):
-        return 1
+        return self.builder.num_ctas
 
     def convert_layout(self, value: Any, layout: Any, assert_trivial: bool = False):
         ty = value.type
@@ -1075,6 +1075,7 @@ def __init__(self) -> None:
         # stale ids from previous simulated launches.
         self._barrier_states: dict[int, _MBarrierState] = {}
         self._pending_clc_barriers: set[int] = set()
+        self.num_ctas = 1
 
     def _barrier_key(self, barrier: Any) -> int:
         handle = getattr(barrier, "handle", None)
@@ -2421,6 +2422,7 @@ def run(self, *args_dev: Any, grid: Any, warmup: bool = False, **kwargs: Any):
 
         client_manager = kwargs.pop("client_manager", None)
         should_patch_lang = kwargs.pop("_patch_lang", True)
+        launch_num_ctas = int(kwargs.get("num_ctas", 1))
         if "num_warps" not in self.arg_names:
             kwargs.pop("num_warps", None)
         if "num_ctas" not in self.arg_names:
@@ -2436,6 +2438,8 @@ def run(self, *args_dev: Any, grid: Any, warmup: bool = False, **kwargs: Any):
         canonical_grid = self._canonical_grid(grid, raw_call_args)
         interpreter_builder.set_grid_dim(*canonical_grid)
         gluon_builder.set_grid_dim(*canonical_grid)
+        previous_num_ctas = gluon_builder.num_ctas
+        gluon_builder.num_ctas = launch_num_ctas
 
         if client_manager is not None:
             for name, arg in raw_call_args.items():
@@ -2465,6 +2469,7 @@ def run_grid() -> None:
                 raise
             raise InterpreterError(repr(exc)) from exc
         finally:
+            gluon_builder.num_ctas = previous_num_ctas
             patch_scope.restore()
 
         grid_helper._restore_args_dev(args_dev, args_hst, kwargs, kwargs_hst)