[Gluon][TTNG] Add async_copy ops including mbarrier arrive op (#7220)

This adds:
- `ttgl.nvidia.ampere.async_copy.async_copy_global_to_shared`
- `ttgl.nvidia.ampere.async_copy.mbarrier_arrive`
- `ttgl.nvidia.ampere.async_copy.commit_group`
- `ttgl.nvidia.ampere.async_copy.wait_group`
- `ttgl.max_constancy`
- `ttgl.max_contiguous`
- `ttgl.multiple_of`
    
Plus adding a new `ttng.async_copy_mbarrier_arrive` op to allow mbarrier
synchronization on `cp.async` ops.
    
The interface on this one is a bit odd since it looks like you can pass
any pointer and mask like a normal load, but it will fail to convert to
llvm if the layout can't be proven to be compatible with `cp.async`.
Hence why I exposed the axis analysis annotation ops.
    
There might be a better interface where you explicitly load chunks of
4, 8, or 16 bytes so we don't rely on axis analysis. This seems okay as
a first draft though.

Author: peterbell10

include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td

@@ -88,7 +88,7 @@ def TTG_AsyncCopyGlobalToLocalOp : TTG_Op<"async_copy_global_to_local", [
   let description = [{
     This operation copies data from global memory to local memory asynchronously.
     This is analogue to tt.load except the data are copied to local memory pointed
-    by by the memory descriptor instead of a distributed tensor. The rest of the
+    to by the memory descriptor instead of a distributed tensor. The rest of the
     operands are the same as tt.load.
   }];
 

include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td

@@ -262,6 +262,15 @@ def TTNG_ArriveBarrierOp : TTNG_Op<"arrive_barrier"> {
   let hasVerifier = 1;
 }
 
+def TTNG_AsyncCopyMbarrierArriveOp : TTNG_Op<"async_copy_mbarrier_arrive"> {
+  let summary = "arrive on mbarrier once all previously issued copies are completed";
+  let arguments = (ins
+    Arg<TTG_MemDescType, "", [MemWrite<SharedMemory>]>:$barrier,
+    UnitAttr:$noIncrement
+  );
+  let assemblyFormat = "$barrier attr-dict `:` qualified(type($barrier))";
+}
+
 
 def TTNG_AsyncTMACopyGlobalToLocalOp : TTNG_Op<"async_tma_copy_global_to_local"> {
   let summary = "copy data based on descriptor from global memory to local memory asynchronously";

python/src/gluon_ir.cc

@@ -279,6 +279,29 @@ void init_gluon_ir(py::module &&m) {
                  blockTy.getShape(), blockTy.getElementType(), layout);
              return triton::TensorDescType::get(ctx, blockTyLayout, isSigned);
            })
+      .def("create_async_copy_global_to_local",
+           [](GluonOpBuilder &self, Value smem, Value pointer, Value mask,
+              tt::CacheModifier cacheModifier,
+              tt::EvictionPolicy evictionPolicy, bool isVolatile) {
+             self.create<ttg::AsyncCopyGlobalToLocalOp>(
+                 pointer, smem, mask, /*other*/ Value{}, cacheModifier,
+                 evictionPolicy, isVolatile);
+           })
+      .def("create_async_copy_mbarrier_arrive",
+           [](GluonOpBuilder &self, Value mbarrier, bool incrementCount) {
+             self.create<ttng::AsyncCopyMbarrierArriveOp>(mbarrier,
+                                                          !incrementCount);
+           })
+      .def("create_async_commit_group",
+           [](GluonOpBuilder &self) {
+             ValueRange tokens;
+             self.create<ttg::AsyncCommitGroupOp>(tokens);
+           })
+      .def("create_async_wait_group",
+           [](GluonOpBuilder &self, int num) {
+             ValueRange tokens;
+             self.create<ttg::AsyncWaitOp>(tokens, num);
+           })
       .def("create_convert_layout",
            [](GluonOpBuilder &self, Type resultTy, Value value) -> Value {
              return self.create<ttg::ConvertLayoutOp>(resultTy, value);

python/test/gluon/test_core.py

@@ -1,9 +1,10 @@
 import torch
 import pytest
 
-from triton._internal_testing import is_cuda
+from triton._internal_testing import is_ampere_or_newer, is_hopper
 from triton.experimental import gluon
 from triton.experimental.gluon import language as ttgl
+from triton.experimental.gluon.language.nvidia.ampere import async_copy, mbarrier
 from triton.experimental.gluon.language.nvidia.hopper import tma
 
 
@@ -45,7 +46,7 @@ def tma_kernel(desc):
     alloc._keep_alive()
 
 
-@pytest.mark.skipif(not is_cuda() or torch.cuda.get_device_capability()[0] < 9, reason="Requires Hopper")
+@pytest.mark.skipif(not is_hopper(), reason="Requires Hopper")
 def test_tma():
     out = torch.ones((16, 16), dtype=torch.float16, device="cuda")
     layout = ttgl.NVMMASharedLayout(
@@ -59,3 +60,36 @@ def test_tma():
     desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(out, [16, 16], layout)
     tma_kernel[(1, )](desc)
     torch.testing.assert_close(out, torch.zeros_like(out))
+
+
+@gluon.jit
+def async_copy_mbarrier_kernel(out, inp, xnumel, XBLOCK: ttgl.constexpr, YBLOCK: ttgl.constexpr):
+    smem = ttgl.allocate_shared_memory(inp.dtype.element_ty, [XBLOCK, YBLOCK],
+                                       ttgl.SwizzledSharedLayout(1, 1, 1, order=[1, 0]))
+    block_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 4], [1, 32], [4, 1], [1, 0])
+    xindex = ttgl.arange(0, XBLOCK, ttgl.SliceLayout(1, block_layout))[:, None]
+    yindex = ttgl.arange(0, YBLOCK, ttgl.SliceLayout(0, block_layout))[None, :]
+    mask = xindex < xnumel
+    async_copy.async_copy_global_to_shared(
+        smem,
+        inp + xindex * YBLOCK + yindex,
+        mask,
+    )
+    mbar = ttgl.allocate_shared_memory(ttgl.int64, [1], mbarrier.MBarrierLayout())
+    mbarrier.init(mbar, count=1)
+    async_copy.mbarrier_arrive(mbar)
+    mbarrier.arrive(mbar)
+    mbarrier.wait(mbar, 0)
+
+    val = smem.load(block_layout)
+    ttgl.store(out + xindex * YBLOCK + yindex, val)
+
+
+@pytest.mark.skipif(not is_ampere_or_newer(), reason="Requires Ampere")
+def test_async_copy_mbarrier():
+    tensor_opts = dict(dtype=torch.float, device="cuda")
+    out = torch.empty((32, 32), **tensor_opts)
+    inp = torch.randn((20, 32), **tensor_opts)
+    async_copy_mbarrier_kernel[(1, )](out, inp, inp.shape[0], XBLOCK=32, YBLOCK=32)
+    torch.testing.assert_close(out[:20], inp)
+    torch.testing.assert_close(out[20:], torch.zeros((12, 32), **tensor_opts))

python/test/gluon/test_frontend.py

@@ -8,11 +8,11 @@
 from triton.experimental import gluon
 from triton.experimental.gluon import language as ttgl
 from triton.experimental.gluon.language.nvidia import blackwell
-from triton.experimental.gluon.language.nvidia.blackwell import mbarrier, tma, TensorMemoryLayout
+from triton.experimental.gluon.language.nvidia.blackwell import mbarrier, tma, TensorMemoryLayout, async_copy
 from triton.experimental.gluon.nvidia.hopper import TensorDescriptor
 from triton._filecheck import filecheck_test, run_parser
 import triton.language as tl
-from triton._internal_testing import is_cuda
+from triton._internal_testing import is_ampere_or_newer, is_blackwell, is_hopper
 from triton.compiler.errors import CompilationError, CompileTimeAssertionFailure
 
 TARGET_PAT = re.compile('ttg.target = "[^"]*"')
@@ -117,8 +117,7 @@ def tensor_memory_kernel(layout: ttgl.constexpr, tmem_layout: ttgl.constexpr):
         buffers.index(i).load(layout)
 
 
-@pytest.mark.skipif(not is_cuda() or torch.cuda.get_device_capability()[0] != 10,
-                    reason="Requires blackwell tensor cores")
+@pytest.mark.skipif(not is_blackwell(), reason="Requires blackwell tensor cores")
 def test_tensor_memory(fresh_knobs):
     knobs.compilation.disable_line_info = True
 
@@ -373,13 +372,13 @@ 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)
+    mbarrier.arrive(bar, count=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")
+@pytest.mark.skipif(not is_hopper(), reason="Requires hopper or newer")
 def test_mbarrier(fresh_knobs):
     knobs.compilation.disable_line_info = True
 
@@ -415,8 +414,7 @@ def tcgen05_mma_kernel(nvmma_layout: ttgl.constexpr, acc_layout: ttgl.constexpr)
     blackwell.tcgen05_mma(a, b, acc)
 
 
-@pytest.mark.skipif(not is_cuda() or torch.cuda.get_device_capability()[0] != 10,
-                    reason="Requires blackwell tensor core")
+@pytest.mark.skipif(not is_blackwell(), reason="Requires blackwell tensor core")
 def test_tcgen05_mma(fresh_knobs):
     knobs.compilation.disable_line_info = True
 
@@ -460,7 +458,7 @@ def async_tma_kernel(input_desc, XBLOCK: ttgl.constexpr):
     tma.store_wait(0)
 
 
-@pytest.mark.skipif(not is_cuda() or torch.cuda.get_device_capability()[0] < 9, reason="TMA requires at least Hopper")
+@pytest.mark.skipif(not is_hopper(), reason="TMA requires at least Hopper")
 def test_async_tma(fresh_knobs):
     knobs.compilation.disable_line_info = True
 
@@ -519,7 +517,7 @@ def async_tma_blackwell_kernel(input_desc, XBLOCK: ttgl.constexpr):
     tma.store_wait(0)
 
 
-@pytest.mark.skipif(not is_cuda() or torch.cuda.get_device_capability()[0] != 10, reason="Requires Blackwell")
+@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
 def test_async_tma_blackwell(fresh_knobs):
     knobs.compilation.disable_line_info = True
 
@@ -955,3 +953,53 @@ def test_inline_asm_elementwise():
     x = ttgl.arange(0, 16, layout)
     # CHECK: elementwise_inline_asm {{.*}} : tensor<16xi32, [[BLOCKED:#.*]]> -> tensor<16xi32, [[BLOCKED]]>
     ttgl.inline_asm_elementwise("mov $0, $0;", "=r,r", [x], dtype=x.dtype, is_pure=True, pack=1)
+
+
+@gluon.jit
+def async_copy_kernel(inp, xnumel, XBLOCK: ttgl.constexpr):
+    smem = ttgl.allocate_shared_memory(inp.dtype.element_ty, [XBLOCK], ttgl.SwizzledSharedLayout(1, 1, 1, order=[0]))
+    block_layout: ttgl.constexpr = ttgl.BlockedLayout([2], [32], [4], [0])
+    xindex = ttgl.arange(0, XBLOCK, block_layout)
+    mask = tl.max_constancy(xindex < xnumel, 2)
+
+    async_copy.async_copy_global_to_shared(smem, inp + xindex, mask)
+    async_copy.async_copy_global_to_shared(smem, inp + xindex, mask, cache_modifier=".ca", eviction_policy="evict_last",
+                                           volatile=True)
+
+    mbar = ttgl.allocate_shared_memory(ttgl.int64, [1], mbarrier.MBarrierLayout())
+    async_copy.mbarrier_arrive(mbar)
+    async_copy.mbarrier_arrive(mbar, increment_count=False)
+    async_copy.commit_group()
+    async_copy.wait_group(0)
+
+
+@pytest.mark.skipif(not is_ampere_or_newer(), reason="Requires ampere")
+def test_async_copy(fresh_knobs):
+    knobs.compilation.disable_line_info = True
+
+    h = async_copy_kernel.warmup(MockTensor(ttgl.float16), xnumel=100, XBLOCK=128, sanitize_overflow=False, grid=(1, ))
+    expecttest.assert_expected_inline(
+        anonymize_ir(h.asm["ttgir"]), """\
+#blocked = #ttg.blocked<{sizePerThread = [2], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
+#loc = loc(unknown)
+#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}>
+#smem = #ttg.shared_memory
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "...", "ttg.threads-per-warp" = 32 : i32} {
+  tt.func public @async_copy_kernel(%arg0: !tt.ptr<f16> {tt.divisibility = 16 : i32} loc(unknown), %arg1: i32 loc(unknown)) attributes {noinline = false} {
+    %0 = ttg.local_alloc : () -> !ttg.memdesc<128xf16, #shared, #smem, mutable> loc(#loc)
+    %1 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32, #blocked> loc(#loc)
+    %2 = tt.splat %arg1 : i32 -> tensor<128xi32, #blocked> loc(#loc)
+    %3 = arith.cmpi slt, %1, %2 {tt.constancy = dense<2> : tensor<1xi32>} : tensor<128xi32, #blocked> loc(#loc)
+    %4 = tt.splat %arg0 : !tt.ptr<f16> -> tensor<128x!tt.ptr<f16>, #blocked> loc(#loc)
+    %5 = tt.addptr %4, %1 : tensor<128x!tt.ptr<f16>, #blocked>, tensor<128xi32, #blocked> loc(#loc)
+    %6 = ttg.async_copy_global_to_local %5, %0 mask %3 : tensor<128x!tt.ptr<f16>, #blocked> -> <128xf16, #shared, #smem, mutable> loc(#loc)
+    %7 = ttg.async_copy_global_to_local %5, %0 mask %3 cacheModifier = ca evictionPolicy = evict_last {isVolatile = true} : tensor<128x!tt.ptr<f16>, #blocked> -> <128xf16, #shared, #smem, mutable> loc(#loc)
+    %8 = ttg.local_alloc : () -> !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc)
+    ttng.async_copy_mbarrier_arrive %8 : !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc)
+    ttng.async_copy_mbarrier_arrive %8 {noIncrement} : !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc)
+    %9 = ttg.async_commit_group  loc(#loc)
+    %10 = ttg.async_wait  {num = 0 : i32} loc(#loc)
+    tt.return loc(#loc)
+  } loc(#loc)
+} loc(#loc)
+""")

python/triton/_internal_testing.py

@@ -38,6 +38,14 @@ def is_cuda():
     return False if target is None else target.backend == "cuda"
 
 
+def is_ampere_or_newer():
+    return is_cuda() and torch.cuda.get_device_capability()[0] >= 8
+
+
+def is_blackwell():
+    return is_cuda() and torch.cuda.get_device_capability()[0] == 10
+
+
 def is_hopper():
     return is_cuda() and torch.cuda.get_device_capability()[0] >= 9
 

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

@@ -44,10 +44,14 @@
 
 _IMPORT_FROM_TRITON: List[str] = [
     "expand_dims",
+    "inline_asm_elementwise",
     "join",
     "load",
     "maximum",
+    "max_constancy",
+    "max_contiguous",
     "minimum",
+    "multiple_of",
     "permute",
     "program_id",
     "reduce",
@@ -58,7 +62,6 @@
     "store",
     "to_tensor",
     "where",
-    "inline_asm_elementwise",
 ]
 
 __all__ = [

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

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

python/triton/experimental/gluon/language/nvidia/ampere/async_copy.py

@@ -0,0 +1,47 @@
+from ..._semantic import _check
+from ..._core import _unwrap_if_constexpr, builtin
+from triton._C.libtriton import ir
+
+__all__ = [
+    "async_copy_global_to_shared",
+    "mbarrier_arrive",
+    "commit_group",
+    "wait_group",
+]
+
+
+@builtin
+def async_copy_global_to_shared(smem, pointer, mask=None, cache_modifier="", eviction_policy="", volatile=False,
+                                _semantic=None):
+    mask = _unwrap_if_constexpr(mask)
+    cache_modifier = _semantic._str_to_load_cache_modifier(cache_modifier)
+    eviction_policy = _semantic._str_to_eviction_policy(eviction_policy)
+    volatile = _unwrap_if_constexpr(volatile)
+    if mask is not None:
+        pointer, mask = _semantic.broadcast_impl_value(pointer, mask)
+    _check(
+        smem.shape == pointer.shape, lambda:
+        f"expected smem shape to match pointer shape but got smem.shape = {smem.shape}, pointer.shape = {pointer.shape}"
+    )
+    mask_handle = mask.handle if mask is not None else ir.value()
+    _semantic.builder.create_async_copy_global_to_local(smem.handle, pointer.handle, mask_handle, cache_modifier,
+                                                        eviction_policy, volatile)
+
+
+@builtin
+def mbarrier_arrive(mbarrier, increment_count=True, _semantic=None):
+    """Arrive on the mbarrier once all outstanding async copies are complete.
+    """
+    increment_count = _unwrap_if_constexpr(increment_count)
+    _semantic.builder.create_async_copy_mbarrier_arrive(mbarrier.handle, increment_count)
+
+
+@builtin
+def commit_group(_semantic=None):
+    _semantic.builder.create_async_commit_group()
+
+
+@builtin
+def wait_group(num_outstanding=0, _semantic=None):
+    num_outstanding = _unwrap_if_constexpr(num_outstanding)
+    _semantic.builder.create_async_wait_group(num_outstanding)

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

@@ -0,0 +1,44 @@
+from triton.experimental.gluon.language._layouts import SwizzledSharedLayout
+from triton.experimental.gluon.language._core import builtin, _unwrap_if_constexpr
+
+__all__ = ["arrive", "init", "invalidate", "MBarrierLayout", "wait"]
+
+
+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, _semantic=None):
+    count = _unwrap_if_constexpr(count)
+    _semantic.builder.create_mbarrier_init(mbarrier.handle, count)
+
+
+@builtin
+def invalidate(mbarrier, _semantic=None):
+    _semantic.builder.create_mbarrier_inval(mbarrier.handle)
+
+
+@builtin
+def wait(mbarrier, phase, pred=True, deps=(), _semantic=None):
+    phase = _semantic.to_tensor(phase)
+    pred = _semantic.to_tensor(pred)
+    deps = [x.handle for x in deps]
+    _semantic.builder.create_mbarrier_wait(mbarrier.handle, phase.handle, pred.handle, deps)
+
+
+@builtin
+def arrive(mbarrier, *, pred=True, _semantic=None):
+    count = 1
+    pred = _semantic.to_tensor(pred)
+    _semantic.builder.create_mbarrier_arrive(mbarrier.handle, count, pred.handle)