[Gluon] Add mbarrier (#6997)

This implements:
- `ttgl.SwizzledSharedLayout`
- `ttgl.nvidia.hopper.mbarrier.MBarrierLayout` (convenience wrapper for `SwizzledSharedLayout`)
- `ttgl.nvidia.hopper.mbarrier.init`
- `ttgl.nvidia.hopper.mbarrier.invalidate`
- `ttgl.nvidia.hopper.mbarrier.expect`
- `ttgl.nvidia.hopper.mbarrier.wait`
- `ttgl.nvidia.hopper.mbarrier.arrive`
    
plus aliases in `ttgl.nvidia.blackwell.mbarrier`
    
Note that I'm keeping this API functional to allow interpreting
any shared allocation as an mbarrier. We can wrap with higher level
APIs at a later date if desired.

Author: peterbell10

python/src/gluon_ir.cc

@@ -80,6 +80,17 @@ void init_gluon_ir(py::module &&m) {
                  ctx, swizzleByteWidth, transposed, elementBitwidth, fp4Padded,
                  ctaLayout);
            })
+      .def("get_swizzled_shared_layout",
+           [](GluonOpBuilder &self, int vec, int perPhase, int maxPhase,
+              std::vector<unsigned> &order, std::vector<unsigned> &ctasPerCga,
+              std::vector<unsigned> &ctaSplitNum,
+              std::vector<unsigned> &ctaOrder) -> Attribute {
+             auto ctx = self.getContext();
+             auto ctaLayout = ttg::CTALayoutAttr::get(ctx, ctasPerCga,
+                                                      ctaSplitNum, ctaOrder);
+             return ttg::SwizzledSharedEncodingAttr::get(
+                 ctx, vec, perPhase, maxPhase, order, ctaLayout);
+           })
       .def("get_tensor_memory_layout",
            [](GluonOpBuilder &self, std::vector<unsigned> &block, bool unpacked,
               std::vector<unsigned> &ctaSplitNum) -> Attribute {
@@ -132,6 +143,27 @@ void init_gluon_ir(py::module &&m) {
               int N) -> Value {
              return self.create<ttng::TMEMSubSliceOp>(resultTy, memDesc, N);
            })
+      .def("create_mbarrier_init",
+           [](GluonOpBuilder &self, Value memDesc, int count) {
+             self.create<ttng::InitBarrierOp>(memDesc, count);
+           })
+      .def("create_mbarrier_inval",
+           [](GluonOpBuilder &self, Value memDesc) {
+             self.create<ttng::InvalBarrierOp>(memDesc);
+           })
+      .def("create_mbarrier_expect",
+           [](GluonOpBuilder &self, Value memDesc, int bytes, Value pred) {
+             self.create<ttng::BarrierExpectOp>(memDesc, bytes, pred);
+           })
+      .def("create_mbarrier_wait",
+           [](GluonOpBuilder &self, Value memDesc, Value phase, Value pred,
+              std::vector<Value> &deps) {
+             self.create<ttng::WaitBarrierOp>(memDesc, phase, pred, deps);
+           })
+      .def("create_mbarrier_arrive",
+           [](GluonOpBuilder &self, Value memDesc, int count, Value pred) {
+             self.create<ttng::ArriveBarrierOp>(memDesc, count, pred);
+           })
       .def("create_warp_return",
            [](GluonOpBuilder &self) -> Operation * {
              return self.create<ttg::WarpReturnOp>();

python/test/gluon/test_frontend.py

@@ -5,6 +5,7 @@
 from triton import knobs
 from triton.experimental import gluon
 from triton.experimental.gluon import language as ttgl
+from triton.experimental.gluon.language.nvidia.blackwell import mbarrier
 from triton._filecheck import filecheck_test
 import triton.language as tl
 from triton._internal_testing import is_cuda
@@ -177,3 +178,42 @@ def test_warp_specialize():
                                 [warp_specialize_worker0, warp_specialize_worker1], [4, 4], [24, 48])
     anchor(a)
     anchor(b)
+
+
+@gluon.jit
+def mbarrier_kernel():
+    bar = ttgl.allocate_shared_memory(ttgl.int64, [1], mbarrier.MBarrierLayout())
+    mbarrier.init(bar, count=1)
+    mbarrier.expect(bar, 4)
+    mbarrier.arrive(bar, 1)
+    phase = 0
+    mbarrier.wait(bar, phase, deps=[bar])
+    mbarrier.invalidate(bar)
+
+
+@pytest.mark.skipif(not is_cuda() or torch.cuda.get_device_capability()[0] < 9, reason="Requires hopper or newer")
+def test_mbarrier(fresh_knobs):
+    knobs.compilation.disable_line_info = True
+
+    h = mbarrier_kernel.warmup(grid=(1, ))
+    expecttest.assert_expected_inline(
+        h.asm["ttgir"], """\
+#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}>
+#smem = #ttg.shared_memory
+module attributes {"ttg.num-warps" = 4 : i32} {
+  tt.func public @mbarrier_kernel() attributes {noinline = false} {
+    %0 = ttg.local_alloc : () -> !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc)
+    ttng.init_barrier %0, 1 : !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc)
+    %true = arith.constant true loc(#loc)
+    ttng.barrier_expect %0, 4, %true : !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc)
+    %true_0 = arith.constant true loc(#loc)
+    ttng.arrive_barrier %0, 1, %true_0 : !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc)
+    %c0_i32 = arith.constant 0 : i32 loc(#loc)
+    %true_1 = arith.constant true loc(#loc)
+    ttng.wait_barrier %0, %c0_i32, %true_1 deps %0 : !ttg.memdesc<1xi64, #shared, #smem, mutable>, !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc)
+    ttng.inval_barrier %0 : !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc)
+    tt.return loc(#loc)
+  } loc(#loc)
+} loc(#loc)
+#loc = loc(unknown)
+""")

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

@@ -2,7 +2,12 @@
 from typing import List, Optional
 from triton.language.core import _unwrap_if_constexpr
 
-__all__ = ["BlockedLayout", "SliceLayout", "NVMMASharedLayout"]
+__all__ = [
+    "BlockedLayout",
+    "SliceLayout",
+    "NVMMASharedLayout",
+    "SwizzledSharedLayout",
+]
 
 
 def _realize_cta_layout(rank, ctas_per_cga, cta_split_num, cta_order):
@@ -123,3 +128,37 @@ def _to_ir(self, builder):
 
     def mangle(self) -> str:
         return f"NVMMA_{self.swizzle_byte_width}_{self.element_bitwidth}_{self.transposed}_{self.fp4_padded}_NVMMA"
+
+
+@dataclass(frozen=True, eq=True)
+class SwizzledSharedLayout(SharedLayout):
+    vec: int
+    per_phase: int
+    max_phase: int
+    order: List[int]
+    ctas_per_cga: Optional[List[int]] = None
+    cta_split_num: Optional[List[int]] = None
+    cta_order: Optional[List[int]] = None
+
+    def __post_init__(self):
+        rank = len(self.order)
+        assert self.ctas_per_cga is None or len(self.ctas_per_cga) == rank
+        assert self.cta_split_num is None or len(self.cta_split_num) == rank
+        assert self.cta_order is None or len(self.cta_order) == rank
+
+    def _to_ir(self, builder):
+        rank = len(self.order)
+        ctas_per_cga, cta_split_num, cta_order = _realize_cta_layout(rank, self.ctas_per_cga, self.cta_split_num,
+                                                                     self.cta_order)
+        return builder.get_swizzled_shared_layout(
+            _unwrap_if_constexpr(self.vec),
+            _unwrap_if_constexpr(self.per_phase),
+            _unwrap_if_constexpr(self.max_phase),
+            self.order,
+            ctas_per_cga,
+            cta_split_num,
+            cta_order,
+        )
+
+    def mangle(self) -> str:
+        return f"NVMMA_{self.swizzle_byte_width}_{self.element_bitwidth}_{self.transposed}_{self.fp4_padded}_NVMMA"

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

@@ -1,3 +1,4 @@
 from . import blackwell
+from . import hopper
 
-__all__ = ["blackwell"]
+__all__ = ["blackwell", "hopper"]

python/triton/experimental/gluon/language/nvidia/blackwell.py renamed to python/triton/experimental/gluon/language/nvidia/blackwell/__init__.py

@@ -5,11 +5,18 @@
 from triton.experimental.gluon.language import _core as ttgl
 from triton.experimental.gluon.language._core import builtin, base_type, base_value, _unwrap_if_constexpr
 
+from ..hopper import mbarrier
+
 if TYPE_CHECKING:
     from triton._C.libtriton.gluon_ir import GluonOpBuilder
     from triton._C.libtriton import gluon_ir as ir
 
-__all__ = ["TensorMemoryLayout", "tensor_memory_descriptor", "allocate_tensor_memory"]
+__all__ = [
+    "TensorMemoryLayout",
+    "tensor_memory_descriptor",
+    "allocate_tensor_memory",
+    "mbarrier",
+]
 
 
 @dataclass(frozen=True, eq=True)

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

@@ -0,0 +1,3 @@
+from . import mbarrier
+
+__all__ = ["mbarrier"]

python/triton/experimental/gluon/language/nvidia/hopper/mbarrier.py

@@ -0,0 +1,52 @@
+from triton.experimental.gluon.language._layouts import SwizzledSharedLayout
+import triton.experimental.gluon.language._core as ttgl
+from triton.experimental.gluon.language._core import builtin, _unwrap_if_constexpr
+
+__all__ = ["MBarrierLayout", "init", "invalidate", "expect", "wait", "arrive"]
+
+
+class MBarrierLayout(SwizzledSharedLayout):
+
+    def __init__(self, ctas_per_cga: int = 1, cta_split_num: int = 1):
+        super().__init__(
+            vec=1,
+            per_phase=1,
+            max_phase=1,
+            order=[0],
+            ctas_per_cga=[ctas_per_cga],
+            cta_split_num=[cta_split_num],
+            cta_order=[0],
+        )
+
+
+@builtin
+def init(mbarrier, count, _builder=None):
+    count = _unwrap_if_constexpr(count)
+    _builder.create_mbarrier_init(mbarrier.handle, count)
+
+
+@builtin
+def invalidate(mbarrier, _builder=None):
+    _builder.create_mbarrier_inval(mbarrier.handle)
+
+
+@builtin
+def expect(mbarrier, bytes, pred=True, _builder=None):
+    bytes = _unwrap_if_constexpr(bytes)
+    pred = ttgl.to_tensor(pred, _builder=_builder)
+    _builder.create_mbarrier_expect(mbarrier.handle, bytes, pred.handle)
+
+
+@builtin
+def wait(mbarrier, phase, pred=True, deps=(), _builder=None):
+    phase = ttgl.to_tensor(phase, _builder=_builder)
+    pred = ttgl.to_tensor(pred, _builder=_builder)
+    deps = [x.handle for x in deps]
+    _builder.create_mbarrier_wait(mbarrier.handle, phase.handle, pred.handle, deps)
+
+
+@builtin
+def arrive(mbarrier, count, pred=True, _builder=None):
+    count = _unwrap_if_constexpr(count)
+    pred = ttgl.to_tensor(pred, _builder=_builder)
+    _builder.create_mbarrier_arrive(mbarrier.handle, count, pred.handle)