[gluon] Basic multi-cta support (#8468)

Currently `num_ctas` is not propagated to the kernel launcher, so just
need to add some minimal metadata.

Author: peterbell10

python/test/gluon/test_core.py

@@ -67,6 +67,18 @@ def test_copy_kernel(layout, XBLOCK):
     torch.testing.assert_close(out, inp)
 
 
+@pytest.mark.skipif(not is_hopper_or_newer(), reason="Requires Hopper")
+def test_copy_kernel_multi_cta():
+    XBLOCK = 2048
+    layout = ttgl.BlockedLayout(size_per_thread=[8], threads_per_warp=[THREADS_PER_WARP], warps_per_cta=[8], order=[0],
+                                ctas_per_cga=[2], cta_split_num=[2])
+
+    inp = torch.randn(XBLOCK * 4 - 7, device="cuda")
+    out = torch.empty_like(inp)
+    copy_kernel[(4, )](out, inp, inp.numel(), XBLOCK, layout, num_warps=layout.warps_per_cta[0], num_ctas=2)
+    torch.testing.assert_close(out, inp)
+
+
 @gluon.jit
 def tma_kernel(desc):
     layout: ttgl.constexpr = ttgl.BlockedLayout([1, 2], [4, 8], [4, 1], [1, 0])

third_party/nvidia/backend/compiler.py

@@ -323,8 +323,7 @@ def make_ttgir(mod, metadata, opt, capability):
 
         pm.run(mod, 'make_ttgir')
         metadata["cluster_dims"] = (cluster_info.clusterDimX, cluster_info.clusterDimY, cluster_info.clusterDimZ)
-        tensordesc_meta = mod.get_tensordesc_metadata()
-        metadata["tensordesc_meta"] = tensordesc_meta
+        metadata["tensordesc_meta"] = mod.get_tensordesc_metadata()
         return mod
 
     def gluon_to_ttgir(self, src, metadata, options, capability):
@@ -341,6 +340,7 @@ def gluon_to_ttgir(self, src, metadata, options, capability):
         passes.ttgpuir.add_combine_tensor_select_and_if(pm)
 
         pm.run(mod, 'gluon_to_ttgir')
+        metadata["cluster_dims"] = (options.num_ctas, 1, 1)
         metadata["tensordesc_meta"] = mod.get_tensordesc_metadata()
         return mod