[GLUON] Implement num_ctas (#8602)

Author: lezcano

python/test/gluon/test_frontend.py

@@ -3006,6 +3006,12 @@ def print_num_warps():
     print("num_warps", num_warps)
 
 
+@gluon.jit
+def print_num_ctas():
+    num_ctas: ttgl.constexpr = ttgl.num_ctas()
+    print("num_ctas", num_ctas)
+
+
 @filecheck_test
 @gluon.jit
 def test_get_num_warps():
@@ -3030,6 +3036,15 @@ def test_get_num_warps():
     ], [1, 2, 8], [24, 24, 24])
 
 
+@filecheck_test
+@gluon.jit
+def test_num_ctas():
+    # CHECK-LABEL: test_num_ctas
+    # CHECK: tt.func private @{{.*}}print_num_ctas
+    # CHECK-NEXT: arith.constant 1 : i32
+    print_num_ctas()
+
+
 def test_mismatch_shape_and_layout_rank():
 
     @gluon.jit

python/triton/experimental/gluon/language/__init__.py

@@ -52,6 +52,7 @@
     full,
     gather,
     num_warps,
+    num_ctas,
     histogram,
     inline_asm_elementwise,
     join,

python/triton/experimental/gluon/language/_core.py

@@ -74,6 +74,7 @@
     "static_range",
     "tuple",
     "tuple_type",
+    "num_ctas",
 ]
 
 T = TypeVar("T")
@@ -525,6 +526,14 @@ def num_warps(_semantic=None, _generator=None):
     return _semantic.num_warps(_generator)
 
 
+@builtin
+def num_ctas(_semantic=None):
+    """
+    Returns the number of CTAs in the current kernel
+    """
+    return _semantic.num_ctas()
+
+
 @builtin
 def thread_barrier(_semantic=None):
     """

python/triton/experimental/gluon/language/_semantic.py

@@ -551,6 +551,9 @@ def warp_specialize(self, functions_and_args, worker_num_warps: Sequence[int], w
             return
         return tuple(unflatten_ir_values(mlir_results, [r.type for r in default_results]))
 
+    def num_ctas(self):
+        return ttgl.constexpr(self.builder.options.num_ctas)
+
     def num_warps(self, generator):
         if generator.caller_context is not None:
             assert isinstance(generator.caller_context, GluonCallerContext)