Repository: facebookexperimental/triton
Pull request #64: [WS][Blackwell]  Fixes thread mapping issue in  llvm codegen
Author: login: htyu
Labels: name: CLA Signed
description: This label is managed by the Meta Open Source bot.
Created: 2025-03-06T19:40:05Z | Merged: 2025-03-28T15:51:21Z

-  Fixes hangs by setting up async id right and the master threading mapping right.
- Fixes incorrect results caused by impropoer thread mapping during TMEM load lowering
    
GEMM kernel runs correctly in one-consumer mode:
    
    python python/gemmbench/gemmbench.py


Diff:
diff --git a/lib/Dialect/TritonNvidiaGPU/Transforms/MMALowering.cpp b/lib/Dialect/TritonNvidiaGPU/Transforms/MMALowering.cpp
index f206c3a366..ca2a8e18ec 100644
--- a/lib/Dialect/TritonNvidiaGPU/Transforms/MMALowering.cpp
+++ b/lib/Dialect/TritonNvidiaGPU/Transforms/MMALowering.cpp
@@ -2,6 +2,7 @@
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 #include "triton/Dialect/Triton/IR/Utility.h"
+#include "triton/Dialect/TritonGPU/Transforms/Utility.h"
 #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h"
 
@@ -38,9 +39,14 @@ class SyncMMALowering : public OpRewritePattern<TCGen5MMAOpTy> {
     MemDescType barrierMemDescType =
         MemDescType::get({1}, rewriter.getI64Type(), barrierEncoding,
                          sharedMemorySpace, /*mutableMemory=*/true);
+
+    PatternRewriterWithAsyncTaskIds taskIdRewriter(rewriter, op);
+
     Value barrierAlloc =
-        rewriter.create<LocalAllocOp>(loc, barrierMemDescType, Value());
-    rewriter.create<InitBarrierOp>(loc, barrierAlloc, 1);
+        taskIdRewriter.create<LocalAllocOp>(loc, barrierMemDescType, Value());
+
+    taskIdRewriter.create<InitBarrierOp>(loc, barrierAlloc, 1);
+
     op.getBarrierMutable().assign(barrierAlloc);
 
     rewriter.setInsertionPointAfter(op);
diff --git a/python/gemmbench/gemmbench.py b/python/gemmbench/gemmbench.py
new file mode 100644
index 0000000000..a1ad4edb4e
--- /dev/null
+++ b/python/gemmbench/gemmbench.py
@@ -0,0 +1,169 @@
+import os
+import sys
+
+import torch
+import triton
+import triton.language as tl
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+import impls
+
+
+if torch.cuda.is_available():
+    from triton._C.libtriton import nvidia
+
+    cublas_workspace = torch.empty(32 * 1024 * 1024, device="cuda", dtype=torch.uint8)
+    cublas = nvidia.cublas.CublasLt(cublas_workspace)
+else:
+    cublas = None
+
+
+def is_cuda():
+    return triton.runtime.driver.active.get_current_target().backend == "cuda"
+
+
+def is_hip_mi200():
+    target = triton.runtime.driver.active.get_current_target()
+    return target.backend == "hip" and target.arch == "gfx90a"
+
+
+def cublas_matmul(a, b):
+    # Check constraints.
+    M, K = a.shape
+    K, N = b.shape
+    dtype = a.dtype
+    c = torch.empty((M, N), device=a.device, dtype=dtype)
+    cublas.matmul(a, b, c)
+    return c
+
+
+def torch_matmul(a, b):
+    M, K = a.shape
+    K, N = b.shape
+    dtype = a.dtype
+    c = torch.empty((M, N), device=a.device, dtype=dtype)
+    c = torch.matmul(a, b.T)
+    return c
+
+
+import torch._dynamo.config
+import torch._inductor.config
+
+torch._dynamo.config.cache_size_limit = 1000000
+# torch._inductor.config.cuda.cutlass_op_allowlist_regex = "(64x256x16.*128x256x64.*cooperative|64x128x16.*128x128x64.*pingpong).*tma"
+# torch._inductor.config.cuda.cutlass_op_denylist_regex = "stream_k"
+
+
+@torch.compile(mode="max-autotune-no-cudagraphs", dynamic=False)
+def inductor_matmul(a, b):
+    return a @ b
+
+
+def aten_matmul(a, b):
+    return a.mm(b)
+
+
+test_impls = [
+    aten_matmul,    
+    impls.matmul_persistent_tma_ws,
+]
+
+impl_map = {fn.__name__: fn for fn in test_impls}
+
+
+def test():
+    torch.manual_seed(0)
+    m = 128
+    n = 128
+    k = 64
+    a = torch.randn((m, k), device="cuda", dtype=torch.float16)
+    b = torch.randn((k, n), device="cuda", dtype=torch.float16)
+    torch_output = torch.matmul(a, b)
+    # Bigger tolerance for AMD MI200 devices.
+    # MI200 devices use reduced precision fp16 and bf16 and flush input and
+    # output denormal values to zero. Detailed info is at: https://pytorch.org/docs/stable/notes/numerical_accuracy.html#reduced-precision-fp16-and-bf16-gemms-and-convolutions-on-amd-instinct-mi200-devices
+    rtol = 1e-2 if is_hip_mi200() else 0
+    for fn in test_impls:
+        triton_output = fn(a, b)
+        torch.cuda.synchronize()
+        if torch.allclose(triton_output, torch_output, atol=1e-2, rtol=rtol):
+            print(f" Torch matches {fn.__name__}")
+        else:
+            print(f" Torch DOES NOT match {fn.__name__}")
+            print("torch output:")
+            print(torch_output)
+            print("triton output:")
+            print(triton_output)
+            #torch.testing.assert_close(triton_output, torch_output, atol=1e-2, rtol=rtol)
+
+
+TORCH_HAS_FP8 = False  # hasattr(torch, "float8_e5m2")
+
+GEMM_SHAPES = os.environ.get("GEMM_SHAPES", "")
+if GEMM_SHAPES == "sweepk":
+    x_vals = [(8192, 8192, i) for i in range(256, 8192 + 1, 256)]
+elif GEMM_SHAPES == "square":
+    x_vals = [(i, i, i) for i in range(256, 4096 + 1, 128)]
+elif GEMM_SHAPES == "llama":
+    KN = [
+        (4096, 12288),
+        (4096, 4096),
+        (4096, 22016),
+        (11008, 4096),
+        (8192, 1280),
+        (1024, 8192),
+        (8192, 7168),
+        (3584, 8192),
+        (16384, 2304),
+        (2048, 16384),
+        (16384, 13312),
+        (6656, 16384),
+    ]
+    x_vals = [(m, n, k) for m in [128, 256, 384, 512] for (k, n) in KN]
+else:
+    # Simple shape with 4 waves over 132 SMs
+    x_vals = [(4 * 11 * 128, 12 * 256, 4096)]
+
+
+configs = []
+for fp8_inputs in [False]:
+    if fp8_inputs and (not TORCH_HAS_FP8 or not is_cuda()):
+        continue
+    configs.append(
+        triton.testing.Benchmark(
+            x_names=["M", "N", "K"],  # Argument names to use as an x-axis for the plot
+            x_vals=x_vals,
+            line_arg="provider",  # Argument name whose value corresponds to a different line in the plot
+            # Possible values for `line_arg`
+            # Don't compare to cublas for fp8 cases as torch.matmul doesn't support fp8 at the moment.
+            line_vals=[fn.__name__ for fn in test_impls],
+            line_names=[fn.__name__ for fn in test_impls],
+            # styles=[("red", "-"), ("green", "-"), ("blue", "-")],
+            ylabel="TFLOPS",  # Label name for the y-axis
+            plot_name="matmul-performance-"
+            + (
+                "fp16" if not fp8_inputs else "fp8"
+            ),  # Name for the plot, used also as a file name for saving the plot.
+            args={"fp8_inputs": fp8_inputs},
+        )
+    )
+
+
+@triton.testing.perf_report(configs)
+def benchmark(M, N, K, provider, fp8_inputs):
+    a = torch.randn((M, K), device="cuda", dtype=torch.float16)
+    b = torch.randn((K, N), device="cuda", dtype=torch.float16)
+    if TORCH_HAS_FP8 and fp8_inputs:
+        a = a.to(torch.float8_e5m2)
+        b = b.T
+        b = b.to(torch.float8_e5m2)
+    quantiles = [0.5, 0.2, 0.8]
+    fn = impl_map[provider]
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(lambda: fn(a, b), quantiles=quantiles)
+    perf = lambda ms: 2 * M * N * K * 1e-12 / (ms * 1e-3)
+    return perf(ms), perf(max_ms), perf(min_ms)
+
+
+test()
+benchmark.run(show_plots=True, print_data=True, save_path=f"./{GEMM_SHAPES}")
diff --git a/python/gemmbench/impls/__init__.py b/python/gemmbench/impls/__init__.py
new file mode 100644
index 0000000000..122d5cc143
--- /dev/null
+++ b/python/gemmbench/impls/__init__.py
@@ -0,0 +1 @@
+from .matmul_persistent_tma_ws import matmul_persistent_tma_ws
diff --git a/python/gemmbench/impls/matmul_persistent_tma_ws.py b/python/gemmbench/impls/matmul_persistent_tma_ws.py
new file mode 100644
index 0000000000..799182fb81
--- /dev/null
+++ b/python/gemmbench/impls/matmul_persistent_tma_ws.py
@@ -0,0 +1,258 @@
+import torch
+import triton
+import triton.language as tl
+
+# check if we have the TMA version in Triton PR #4498 (https://github.com/triton-lang/triton/pull/4498).
+HAS_TMA_DESC = "nv_tma_desc_type" in dir(tl)
+
+if HAS_TMA_DESC:
+    print(
+        "TMA benchmarks will be running with experimental grid constant TMA descriptor.",
+    )
+else:
+    print(
+        "TMA benchmarks will be running without grid constant TMA descriptor.",
+    )
+
+
+class TmaAutoTuneHelper:
+
+    # duck typing wrapper to implement the same interface as TmaDescKernelParam in Triton PR #4498
+    class KernelParamWrapper:
+        def __init__(self, desc):
+            self.desc = desc
+
+        def tma_desc_cpu_ptr(self):
+            return self.desc.data_ptr()
+
+    TMA_SIZE = 128
+
+    def __init__(self):
+        self.fill_1d_tma_descriptor_inner = (
+            triton.runtime.driver.active.utils.fill_1d_tma_descriptor
+        )
+        self.fill_2d_tma_descriptor_inner = (
+            triton.runtime.driver.active.utils.fill_2d_tma_descriptor
+        )
+        if HAS_TMA_DESC:
+            self.descriptors = {}
+        else:
+            self.cuda_descriptors = {}
+
+    # Call this method outside of the lambda function for grid size
+    def init_tma_descriptor(self, name):
+        if HAS_TMA_DESC:
+            self.descriptors[name] = torch.empty(
+                TmaAutoTuneHelper.TMA_SIZE, device="cpu", dtype=torch.int8
+            )
+        else:
+            self.cuda_descriptors[name] = torch.empty(
+                TmaAutoTuneHelper.TMA_SIZE, device="cuda", dtype=torch.int8
+            )
+
+    # Call this method inside the lambda function for grid size
+    def fill_1d_tma_descriptor(self, name, ptr, dim, block_dim, element_size):
+        if HAS_TMA_DESC:
+            desc_x = self.descriptors[name]
+            assert desc_x.data_ptr() % 64 == 0
+            self.fill_1d_tma_descriptor_inner(
+                ptr, dim, block_dim, element_size, desc_x.data_ptr()
+            )
+        else:
+            desc_x = self.cuda_descriptors[name]
+            buf_x = torch.empty_like(desc_x, device="cpu", pin_memory=True)
+            self.fill_1d_tma_descriptor_inner(
+                ptr, dim, block_dim, element_size, buf_x.data_ptr()
+            )
+            desc_x.copy_(buf_x, non_blocking=True)
+
+    # Call this method inside the lambda function for grid size
+    def fill_2d_tma_descriptor(
+        self, name, ptr, dim1, dim0, block_dim1, block_dim0, element_size
+    ):
+        if HAS_TMA_DESC:
+            desc_x = self.descriptors[name]
+            assert desc_x.data_ptr() % 64 == 0
+            self.fill_2d_tma_descriptor_inner(
+                ptr, dim1, dim0, block_dim1, block_dim0, element_size, desc_x.data_ptr()
+            )
+        else:
+            desc_x = self.cuda_descriptors[name]
+            buf_x = torch.empty_like(desc_x, device="cpu", pin_memory=True)
+            self.fill_2d_tma_descriptor_inner(
+                ptr, dim1, dim0, block_dim1, block_dim0, element_size, buf_x.data_ptr()
+            )
+            desc_x.copy_(buf_x, non_blocking=True)
+
+    def get_tma_descriptor_kernel_param(self, name):
+        if HAS_TMA_DESC:
+            assert self.descriptors[name] is not None
+            return self.KernelParamWrapper(self.descriptors[name])
+        else:
+            assert self.cuda_descriptors[name] is not None
+            return self.cuda_descriptors[name]
+
+
+@triton.autotune(
+    configs=[
+        triton.Config(
+            {
+                "BLOCK_SIZE_M": 128,
+                "BLOCK_SIZE_N": 128,
+                "BLOCK_SIZE_K": 64,
+                "GROUP_SIZE_M": 8,
+            },
+            num_stages=2,
+            num_warps=4,
+            num_consumer_groups=1,
+            num_buffers_warp_spec=3,
+        ),
+    ],
+    key=["M", "N", "K"],
+)
+@triton.jit
+def matmul_persistent_tma_ws_kernel(
+    a_ptr,
+    b_ptr,
+    c_ptr,
+    M,
+    N,
+    K,
+    stride_am,
+    stride_ak,  #
+    stride_bk,
+    stride_bn,  #
+    stride_cm,
+    stride_cn,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,  #
+    GROUP_SIZE_M: tl.constexpr,  #
+):
+    """Kernel for computing the matmul C = A x B.
+    A has shape (M, K), B has shape (K, N) and C has shape (M, N)
+    """
+
+    num_tiles = tl.cdiv(M, BLOCK_SIZE_M) * tl.cdiv(N, BLOCK_SIZE_N)
+    for pid in range(tl.program_id(0), num_tiles, tl.num_programs(0)):
+        num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+        num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+        num_pid_in_group = GROUP_SIZE_M * num_pid_n
+        group_id = pid // num_pid_in_group
+        first_pid_m = group_id * GROUP_SIZE_M
+        group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+        pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+        pid_n = (pid % num_pid_in_group) // group_size_m
+
+        # ----------------------------------------------------------
+        # Create pointers for the first blocks of A and B.
+        # We will advance this pointer as we move in the K direction
+        # and accumulate
+        # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers
+        # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers
+        # See above `Pointer Arithmetic` section for details
+        offs_am = pid_m * BLOCK_SIZE_M
+        offs_bn = pid_n * BLOCK_SIZE_N
+        offs_k = 0
+
+        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+        for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+            with tl.async_task([0]):
+                a = tl._experimental_descriptor_load(
+                    a_ptr,
+                    [offs_am, offs_k],
+                    [BLOCK_SIZE_M, BLOCK_SIZE_K],
+                    tl.float16,
+                )
+                b = tl._experimental_descriptor_load(
+                    b_ptr, [offs_k, offs_bn], [BLOCK_SIZE_K, BLOCK_SIZE_N], tl.float16
+                )
+
+            with tl.async_task([1]):
+                accumulator = tl.dot(a, b, accumulator)
+            offs_k += BLOCK_SIZE_K
+
+        c = accumulator.to(tl.float16)
+
+        with tl.async_task([1]):
+            tl._experimental_descriptor_store(c_ptr, c, [offs_am, offs_bn])
+
+
+# %%
+# We can now create a convenience wrapper function that only takes two input tensors,
+# and (1) checks any shape constraint; (2) allocates the output; (3) launches the above kernel.
+
+
+def matmul_persistent_tma_ws(a, b):
+    # Check constraints.
+    assert a.shape[1] == b.shape[0], "Incompatible dimensions"
+    assert a.dtype == b.dtype, "Incompatible dtypes"
+    NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count
+    M, K = a.shape
+    K, N = b.shape
+    dtype = a.dtype
+    # Allocates output.
+    c = torch.empty((M, N), device=a.device, dtype=dtype)
+
+    desc_helper = TmaAutoTuneHelper()
+    desc_helper.init_tma_descriptor("a")
+    desc_helper.init_tma_descriptor("b")
+    desc_helper.init_tma_descriptor("c")
+
+    def grid(META):
+        nonlocal desc_helper
+        desc_helper.fill_2d_tma_descriptor(
+            "a",
+            a.data_ptr(),
+            M,
+            K,
+            META["BLOCK_SIZE_M"],
+            META["BLOCK_SIZE_K"],
+            a.element_size(),
+        )
+
+        desc_helper.fill_2d_tma_descriptor(
+            "b",
+            b.data_ptr(),
+            K,
+            N,
+            META["BLOCK_SIZE_K"],
+            META["BLOCK_SIZE_N"],
+            b.element_size(),
+        )
+        desc_helper.fill_2d_tma_descriptor(
+            "c",
+            c.data_ptr(),
+            M,
+            N,
+            META["BLOCK_SIZE_M"],
+            META["BLOCK_SIZE_N"],
+            c.element_size(),
+        )
+        return (
+            min(
+                NUM_SMS,
+                triton.cdiv(M, META["BLOCK_SIZE_M"])
+                * triton.cdiv(N, META["BLOCK_SIZE_N"]),
+            ),
+        )
+
+    desc_a = desc_helper.get_tma_descriptor_kernel_param("a")
+    desc_b = desc_helper.get_tma_descriptor_kernel_param("b")
+    desc_c = desc_helper.get_tma_descriptor_kernel_param("c")
+    matmul_persistent_tma_ws_kernel[grid](
+        desc_a,
+        desc_b,
+        desc_c,  #
+        M,
+        N,
+        K,  #
+        a.stride(0),
+        a.stride(1),  #
+        b.stride(0),
+        b.stride(1),  #
+        c.stride(0),
+        c.stride(1),  #
+    )
+    return c
diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py
index 6db76a3524..fc4bad74e3 100644
--- a/third_party/nvidia/backend/compiler.py
+++ b/third_party/nvidia/backend/compiler.py
@@ -324,6 +324,8 @@ def make_llir(self, src, metadata, options, capability):
         nvidia.passes.ttnvgpuir.add_lower_mma(pm)
         passes.ttgpuir.add_combine_tensor_select_and_if(pm)
         passes.ttgpuir.add_allocate_warp_groups(pm)
+        if capability // 10 >= 9:
+            passes.ttgpuir.add_ws_canonicalization(pm, options.num_consumer_groups)
         passes.convert.add_scf_to_cf(pm)
         passes.ttgpuir.add_allocate_shared_memory(pm)
         nvidia.passes.ttnvgpuir.add_allocate_tensor_memory(pm)
diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv5.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv5.cpp
index e5ec0df36f..9bc8662c58 100644
--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv5.cpp
+++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv5.cpp
@@ -310,7 +310,16 @@ void convertDot(const LLVMTypeConverter *typeConverter,
   // Only run mma on one thread. We currently use elect as ptxas is not able to
   // detect that tid.x == 0 is true only for 1 thread.
   Value warpId = rewriter.create<nvgpu::WarpIdOp>(loc);
-  Value wapr0 = tb.icmp_eq(warpId, tb.i32_val(0));
+  auto asyncTaskIds = getAsyncTaskIds(op);
+  int executingWarpId = 0;
+  if (!asyncTaskIds.empty()) {
+    assert(asyncTaskIds.size() == 1 && "only support single async task");
+    auto mod = op->getParentOfType<ModuleOp>();
+    int numWarps = triton::gpu::lookupNumWarps(mod);
+    executingWarpId = asyncTaskIds[0] * numWarps;
+  }
+
+  Value wapr0 = tb.icmp_eq(warpId, tb.i32_val(executingWarpId));
   if (twoCTAs) {
     // TODO: we have to sync the two CTAs because we currently don't use remove
     // barriers for the copies.
diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TensorMemoryToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TensorMemoryToLLVM.cpp
index e9e8909d52..bdf6c6192f 100644
--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TensorMemoryToLLVM.cpp
+++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TensorMemoryToLLVM.cpp
@@ -217,7 +217,7 @@ TMemRuntimeInfo getTMemRuntimeInfo(Operation *op, RankedTensorType tensorType,
 }
 
 void calculateAddressAndEmitTmemMessage(
-    Location loc, Value baseAddress, const TMemRuntimeInfo &info,
+    Location loc, Operation *op, Value baseAddress, const TMemRuntimeInfo &info,
     const TMemMessageTraits &message, ConversionPatternRewriter &rewriter,
     const std::function<void(Value, int, bool, int, bool)> &createMemoryOp) {
 
@@ -226,6 +226,15 @@ void calculateAddressAndEmitTmemMessage(
   Value warpIdInGroup = b.urem(warpId, b.i32_val(4));
   Value warpGroupId = b.udiv(warpId, b.i32_val(4));
 
+  auto asyncTaskIds = getAsyncTaskIds(op);
+  if (!asyncTaskIds.empty()) {
+    // numWarpGroups is for the current groups that do the load. In WS mode, we
+    // assume only one warp group does TMEMLoad.
+    assert(asyncTaskIds.size() == 1 &&
+           "only support TMEM load in single async task");
+    warpGroupId = b.sub(warpGroupId, b.i32_val(asyncTaskIds[0]));
+  }
+
   for (int block = 0; block < info.numBlocks; block += info.numWarpGroups) {
     Value address = b.ptrtoint(i32_ty, baseAddress);
     Value blockId =
@@ -349,7 +358,7 @@ static void lowerStoreToTensorMemory(Location loc, Operation *op, Value src,
   const TMemMessageTraits message = selectTMemMessage(info);
   int regIdx = 0;
   calculateAddressAndEmitTmemMessage(
-      loc, tmemBase, info, message, rewriter,
+      loc, op, tmemBase, info, message, rewriter,
       [&](Value startAddress, int secondHalfColOffset, bool unpackedb16,
           int regsPerMsg, bool useStridedMessage) {
         SmallVector<Value> srcValuesSlice(srcValues.begin() + regIdx,
@@ -364,7 +373,7 @@ static void lowerStoreToTensorMemory(Location loc, Operation *op, Value src,
 
   // Emit a barrier to ensure all threads have finished writing to tensor memory
   // before any use of the tensor memory.
-  b.barrier();
+  insertBarrier(rewriter, op);
 }
 
 struct TensorMemoryAllocOpConversion
@@ -487,7 +496,7 @@ struct TensorMemoryLoadOpConversion
     const TMemMessageTraits message = selectTMemMessage(info);
     SmallVector<Value> resultVals;
     calculateAddressAndEmitTmemMessage(
-        loc, tmemBase, info, message, rewriter,
+        loc, op, tmemBase, info, message, rewriter,
         [&](Value startAddress, int secondHalfColOffset, bool unpackedb16,
             int regsPerMessage, bool useStridedMessage) {
           Value packedValues = createTensorMemoryLoad(


Reviews:
user: login: manman-ren
body: Forgot to submit the review :]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T19:06:32Z
user: login: htyu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T22:50:23Z
user: login: manman-ren
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-28T01:01:31Z

Review comments:
diff_hunk: @@ -92,6 +92,16 @@ struct PipelinePass : public impl::TritonGPUPipelineBase<PipelinePass> {
 
   void runOnOperation() override {
     ModuleOp moduleOp = getOperation();
+    unsigned loopCnt = 0;
+    getOperation()->walk([&](scf::ForOp forOp) {
+      // Bail out for loops with num_stage <= 1.
+      if (getNumStagesOrDefault(forOp) > 1)
+        loopCnt++;
+    });
+
+    if (loopCnt == 0)
+      return;
+
path: lib/Dialect/TritonGPU/Transforms/Pipeliner/SoftwarePipeliner.cpp
user: login: manman-ren
body: This doesn't look specific to warp spec. I wonder why we need it when warp spec is on.
created_at: 2025-03-25T22:57:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,169 @@
+import os
+import sys
+
+import torch
+import triton
+import triton.language as tl
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
path: python/gemmbench/gemmbench.py
user: login: manman-ren
body: Is this imported from ws branch?
created_at: 2025-03-25T22:57:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2910,7 +2910,7 @@ class async_task:
     Context manager to run code fragments asynchronously.
     """
     def __init__(self, task_ids, _builder=None):
-        self.task_ids = task_ids
+        self.task_ids = list({_constexpr_to_value(tid) for tid in task_ids})
         self.builder = _builder
path: python/triton/language/core.py
user: login: manman-ren
body: Curious on why we need this for Blackwell.
created_at: 2025-03-25T22:58:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -324,6 +324,8 @@ def make_llir(self, src, metadata, options, capability):
         nvidia.passes.ttnvgpuir.add_lower_mma(pm)
         passes.ttgpuir.add_combine_tensor_select_and_if(pm)
         passes.ttgpuir.add_allocate_warp_groups(pm)
+        if capability // 10 >= 9:
+            passes.ttgpuir.add_ws_canonicalization(pm, options.num_consumer_groups)
path: third_party/nvidia/backend/compiler.py
user: login: manman-ren
body: Is this ported from ws?
created_at: 2025-03-25T22:59:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -226,6 +226,15 @@ void calculateAddressAndEmitTmemMessage(
   Value warpIdInGroup = b.urem(warpId, b.i32_val(4));
   Value warpGroupId = b.udiv(warpId, b.i32_val(4));
 
+  auto asyncTaskIds = getAsyncTaskIds(op);
+  if (!asyncTaskIds.empty()) {
+    // numWarpGroups is for the current groups that do the load. In WS mode, we
+    // assume only one warp group does TMEMLoad.
+    assert(asyncTaskIds.size() == 1 &&
+           "only support TMEM load in single async task");
+    warpGroupId = b.sub(warpGroupId, b.i32_val(asyncTaskIds[0]));
path: third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TensorMemoryToLLVM.cpp
user: login: manman-ren
body: Not sure why we need to do "wgId - taskId", so if task 1 wants to perform tmem load, we will use "wgId - 1", don't quite get why.
created_at: 2025-03-25T23:07:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,169 @@
+import os
+import sys
+
+import torch
+import triton
+import triton.language as tl
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
path: python/gemmbench/gemmbench.py
user: login: htyu
body: Yeah, I found that easier for testing. WDYT?
created_at: 2025-03-27T22:43:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -92,6 +92,16 @@ struct PipelinePass : public impl::TritonGPUPipelineBase<PipelinePass> {
 
   void runOnOperation() override {
     ModuleOp moduleOp = getOperation();
+    unsigned loopCnt = 0;
+    getOperation()->walk([&](scf::ForOp forOp) {
+      // Bail out for loops with num_stage <= 1.
+      if (getNumStagesOrDefault(forOp) > 1)
+        loopCnt++;
+    });
+
+    if (loopCnt == 0)
+      return;
+
path: lib/Dialect/TritonGPU/Transforms/Pipeliner/SoftwarePipeliner.cpp
user: login: htyu
body: WS created a loop pattern that failed tcgen5 pipeliner, even under num_stage=0. But now it seems that the pattern is no longer generated, so the change is not needed.
created_at: 2025-03-27T22:46:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -324,6 +324,8 @@ def make_llir(self, src, metadata, options, capability):
         nvidia.passes.ttnvgpuir.add_lower_mma(pm)
         passes.ttgpuir.add_combine_tensor_select_and_if(pm)
         passes.ttgpuir.add_allocate_warp_groups(pm)
+        if capability // 10 >= 9:
+            passes.ttgpuir.add_ws_canonicalization(pm, options.num_consumer_groups)
path: third_party/nvidia/backend/compiler.py
user: login: htyu
body: Actually no. On Blackwell the `add_allocate_warp_groups` pass could introduce new code that misses task id.
created_at: 2025-03-27T22:48:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -226,6 +226,15 @@ void calculateAddressAndEmitTmemMessage(
   Value warpIdInGroup = b.urem(warpId, b.i32_val(4));
   Value warpGroupId = b.udiv(warpId, b.i32_val(4));
 
+  auto asyncTaskIds = getAsyncTaskIds(op);
+  if (!asyncTaskIds.empty()) {
+    // numWarpGroups is for the current groups that do the load. In WS mode, we
+    // assume only one warp group does TMEMLoad.
+    assert(asyncTaskIds.size() == 1 &&
+           "only support TMEM load in single async task");
+    warpGroupId = b.sub(warpGroupId, b.i32_val(asyncTaskIds[0]));
path: third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TensorMemoryToLLVM.cpp
user: login: htyu
body: It's because the following logics always computes tmem load offset based on starting warp id is zero.
created_at: 2025-03-27T22:49:43Z
author_association: CONTRIBUTOR

PR comments:
user: login: masahi
created_at: 2025-03-06T20:10:04Z
body: I think your `MBarrierArriveOp` after MMA (`consumer_release` on the TMEM full barrier) needs to lower to `tcgen05.commit`. Its similar to the "track async op" semantics https://github.com/facebookexperimental/triton/blob/ws-main/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/BarrierOpToLLVM.cpp#L86-L87, but depending on which barrier it is associated with and the target arch, it might need to lower to `tcgen05.commit`.
author_association: CONTRIBUTOR
user: login: masahi
created_at: 2025-03-06T21:19:19Z
body: We checked the TTGIR (thanks @mbrookhart). It shows that `tc_gen5_mma` is lowered as a blocking operation, and mbar_arrive on the empty barrier follows it. So you don't have the barrier I mentioned above.

Instead, you should give the TMA empty barrier to the barrier parameter of `tc_gen5_mma`. When `tc_gen5_mma` is lowered to PTX, Triton generates `tcgen05.commit` on the provided barrier, which does arrive on the barrier after MMA completes. This is how MMA becomes async and synchronizes with TMA. Then we need synchronization after MMA for one tile completes. That's where my comment above becomes relevant. 
author_association: CONTRIBUTOR
user: login: masahi
created_at: 2025-03-06T21:39:13Z
body: There are several other issues in TTGIR, but the first thing I would address is the two MMA groups structure which mirrors the Hopper one but not really useful for Blackwell. Blackwell MMA is issued by a single thread, so we don't need two warp groups. ~Actually I'm not sure if it is even valid to issues multiple MMAs from two warps in a CTA - probably not.~ (UPDATE: apparently it's ok)

To be clear, does the non-cooperative one work?
author_association: CONTRIBUTOR
user: login: htyu
created_at: 2025-03-10T16:55:18Z
body: > To be clear, does the non-cooperative one work?

It works in terms of non-deadlocking. Perf is getting closer to cublas but has a numeric issue. It looks like the tensor memory load always return 0. I compared the ttgir with the one from non-WS but didn't find anything useful. Wondering if you have help to take a look didn't see anything obviously wrong. Thanks!


TTGIR: https://gist.github.com/htyu/f73b6c25f22f40da84d5b5226e5597cd 
PTX: https://gist.github.com/htyu/d8a2e34a00457b655f14981eb5b12335
author_association: CONTRIBUTOR
user: login: masahi
created_at: 2025-03-10T19:55:11Z
body: I'd watch out for more thread / warp id related issues, see for example

https://github.com/facebookexperimental/triton/blob/3f0bd100bc71c7ba90e2544bdf17b4c21413754d/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TensorMemoryToLLVM.cpp#L70
https://github.com/facebookexperimental/triton/blob/3f0bd100bc71c7ba90e2544bdf17b4c21413754d/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TensorMemoryToLLVM.cpp#L107
author_association: CONTRIBUTOR
user: login: htyu
created_at: 2025-03-13T23:22:23Z
body: Big thanks to @mbrookhart for root-causing the tmem load issue that caused the output error!
author_association: CONTRIBUTOR