Merge commit 'e87fdd9ac41f136619a76ee22cc3d4e926e319de'

Author: whitneywhtsang

lib/Analysis/AxisInfo.cpp

@@ -1109,10 +1109,7 @@ void AxisInfoAnalysis::visitForOpInductionVar(
   AxisInfo::DimVectorT knownContiguity(1, 1);
   AxisInfo::DimVectorT knownDivisibility(1, 1);
   AxisInfo::DimVectorT knownConstancy(1, 1);
-  auto lbDivisibility = lb.getDivisibility();
-  auto stepDivisibility = step.getDivisibility();
-  if (!lbDivisibility.empty() && !stepDivisibility.empty())
-    knownDivisibility[0] = gcd(lbDivisibility[0], stepDivisibility[0]);
+  knownDivisibility[0] = gcd(lb.getDivisibility(0), step.getDivisibility(0));
   auto inductionVar =
       AxisInfo(knownContiguity, knownDivisibility, knownConstancy);
   (void)argLattices[0]->join(inductionVar);

lib/Analysis/Utility.cpp

@@ -1167,63 +1167,10 @@ SetVector<Operation *> multiRootGetSlice(Operation *op,
   return multiRootTopologicalSort(slice);
 }
 
-namespace {
-// Copied from TestDeadCodeAnalysis.cpp, because some dead code analysis
-// interacts with constant propagation, but SparseConstantPropagation
-// doesn't seem to be sufficient.
-class ConstantAnalysis : public DataFlowAnalysis {
-public:
-  using DataFlowAnalysis::DataFlowAnalysis;
-
-  LogicalResult initialize(Operation *top) override {
-    WalkResult result = top->walk([&](Operation *op) {
-      ProgramPoint programPoint(op);
-      if (failed(visit(&programPoint)))
-        return WalkResult::interrupt();
-      return WalkResult::advance();
-    });
-    return success(!result.wasInterrupted());
-  }
-
-  LogicalResult visit(ProgramPoint *point) override {
-    Operation *op = point->getOperation();
-    Attribute value;
-    if (matchPattern(op, m_Constant(&value))) {
-      auto *constant = getOrCreate<dataflow::Lattice<dataflow::ConstantValue>>(
-          op->getResult(0));
-      propagateIfChanged(constant, constant->join(dataflow::ConstantValue(
-                                       value, op->getDialect())));
-      return success();
-    }
-    // Dead code analysis requires every operands has initialized ConstantValue
-    // state before it is visited.
-    // https://github.com/llvm/llvm-project/blob/2ec1aba2b69faa1de5f71832a48e25aa3b5d5314/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp#L322
-    // That's why we need to set all operands to unknown constants.
-    setAllToUnknownConstants(op->getResults());
-    for (Region &region : op->getRegions()) {
-      for (Block &block : region.getBlocks())
-        setAllToUnknownConstants(block.getArguments());
-    }
-    return success();
-  }
-
-private:
-  /// Set all given values as not constants.
-  void setAllToUnknownConstants(ValueRange values) {
-    dataflow::ConstantValue unknownConstant(nullptr, nullptr);
-    for (Value value : values) {
-      auto *constant =
-          getOrCreate<dataflow::Lattice<dataflow::ConstantValue>>(value);
-      propagateIfChanged(constant, constant->join(unknownConstant));
-    }
-  }
-};
-} // namespace
-
 std::unique_ptr<DataFlowSolver> createDataFlowSolver() {
   auto solver = std::make_unique<DataFlowSolver>();
   solver->load<dataflow::DeadCodeAnalysis>();
-  solver->load<ConstantAnalysis>();
+  solver->load<dataflow::SparseConstantPropagation>();
   return solver;
 }
 

lib/Dialect/TritonGPU/IR/Dialect.cpp

@@ -2754,7 +2754,8 @@ struct TritonGPUInferLayoutInterface
     auto mmaRetEncoding = mlir::dyn_cast<NvidiaMmaEncodingAttr>(retEncoding);
     if (mmaRetEncoding && mmaRetEncoding.isHopper()) {
       auto dotOpEnc = mlir::dyn_cast<DotOperandEncodingAttr>(operandEncoding);
-      if (!mlir::isa<NVMMASharedEncodingAttr>(operandEncoding) &&
+      if (!mlir::isa<NVMMASharedEncodingAttr, SharedLinearEncodingAttr>(
+              operandEncoding) &&
           !(opIdx == 0 && dotOpEnc && dotOpEnc.getOpIdx() == 0 &&
             mlir::isa<NvidiaMmaEncodingAttr>(dotOpEnc.getParent()))) {
         return emitOptionalError(

lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp

@@ -50,7 +50,7 @@ LogicalResult WarpGroupDotOp::inferReturnTypes(
 
   // verify encodings
   auto aEnc = cast<TensorOrMemDesc>(operands[0].getType()).getEncoding();
-  auto bEnc = cast<TensorOrMemDesc>(operands[1].getType()).getEncoding();
+  auto bEnc = cast<MemDescType>(operands[1].getType()).getEncoding();
   auto retEnc = accTy.getEncoding();
   if (aEnc) {
     assert(bEnc);
@@ -70,10 +70,11 @@ LogicalResult WarpGroupDotOp::verify() {
   if (!nvmmaEnc || !nvmmaEnc.isHopper())
     return emitOpError("WGMMA result layout must be Hopper NVMMA");
 
-  if (!isa<NVMMASharedEncodingAttr, DotOperandEncodingAttr>(
-          getA().getType().getEncoding()))
+  if (!isa<NVMMASharedEncodingAttr, DotOperandEncodingAttr,
+           SharedLinearEncodingAttr>(getA().getType().getEncoding()))
     return emitOpError("WGMMA A operand must have NVMMA shared or dot layout");
-  if (!isa<NVMMASharedEncodingAttr>(getB().getType().getEncoding()))
+  if (!isa<NVMMASharedEncodingAttr, SharedLinearEncodingAttr>(
+          getB().getType().getEncoding()))
     return emitOpError("WGMMA B operand must have NVMMA shared layout");
 
   auto numWarps = gpu::lookupNumWarps(getOperation());

python/src/gluon_ir.cc

@@ -669,6 +669,18 @@ void init_gluon_ir(py::module &&m) {
                                             pred, two_ctas, mbarriers,
                                             mbarrier_preds);
            })
+      .def("create_tcgen05_mma_scaled",
+           [](GluonOpBuilder &self, Value a, Value b, Value acc, Value aScale,
+              Value bScale, tt::ScaleDotElemType aType,
+              tt::ScaleDotElemType bType, Value useAcc, Value pred,
+              std::vector<Value> &mbarriers,
+              std::vector<Value> &mbarrier_preds) {
+             Value accDep;
+             auto tokType = self.getBuilder().getType<ttg::AsyncTokenType>();
+             self.create<ttng::TCGen5MMAScaledOp>(
+                 tokType, a, b, acc, accDep, aScale, bScale, aType, bType,
+                 useAcc, pred, mbarriers, mbarrier_preds);
+           })
       .def("create_tcgen05_commit",
            [](GluonOpBuilder &self, Value &barrier) {
              self.create<ttng::TCGen5CommitOp>(barrier);

python/src/ir.cc

@@ -455,6 +455,8 @@ void init_triton_ir(py::module &&m) {
              auto loc = UnknownLoc::get(ty.getContext());
              self.addArgument(ty, loc);
            })
+      .def("add_argument_at", [](Block &self, Type ty,
+                                 Location loc) { self.addArgument(ty, loc); })
       .def("get_num_arguments", &Block::getNumArguments)
       .def("get_argument", &Block::getArgument)
       .def("dump", &Block::dump)

python/test/gluon/test_core.py

@@ -1,4 +1,5 @@
 import torch
+import math
 import pytest
 import re
 from itertools import product
@@ -126,9 +127,9 @@ def test_async_copy_mbarrier(device):
 
 
 @gluon.jit
-def warpgroup_mma_kernel(a, b, out, M: ttgl.constexpr, N: ttgl.constexpr, K: ttgl.constexpr,
-                         block_layout: ttgl.constexpr, mma_layout: ttgl.constexpr, shared_layout_a: ttgl.constexpr,
-                         shared_layout_b: ttgl.constexpr, acc_dtype: ttgl.constexpr, ASYNC: ttgl.constexpr):
+def mma_kernel(a, b, out, M: ttgl.constexpr, N: ttgl.constexpr, K: ttgl.constexpr, block_layout: ttgl.constexpr,
+               mma_layout: ttgl.constexpr, shared_layout_a: ttgl.constexpr, shared_layout_b: ttgl.constexpr,
+               acc_dtype: ttgl.constexpr, ASYNC: ttgl.constexpr, USE_TCGEN05: ttgl.constexpr):
     a_offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, block_layout))[:, None]
     a_offs_k = ttgl.arange(0, K, layout=ttgl.SliceLayout(0, block_layout))[None, :]
     b_offs_k = ttgl.arange(0, K, layout=ttgl.SliceLayout(1, block_layout))[:, None]
@@ -143,14 +144,37 @@ def warpgroup_mma_kernel(a, b, out, M: ttgl.constexpr, N: ttgl.constexpr, K: ttg
 
     smem_a = ttgl.allocate_shared_memory(operand_dtype, [M, K], shared_layout_a, a_tile)
     smem_b = ttgl.allocate_shared_memory(operand_dtype, [K, N], shared_layout_b, b_tile)
-
     fence_async_shared()
 
-    acc = ttgl.zeros([M, N], dtype=acc_dtype, layout=mma_layout)
-    acc = hopper.warpgroup_mma(smem_a, smem_b, acc, is_async=ASYNC)
+    if USE_TCGEN05:
+        tmem_layout: ttgl.constexpr = TensorMemoryLayout((M, N), col_stride=32 // acc_dtype.primitive_bitwidth)
+
+        num_warps: ttgl.constexpr = ttgl.num_warps()
+        tmem_reg_layout: ttgl.constexpr = get_tmem_32x32b_reg_layout(
+            M=M,
+            N=N,
+            shape=[M, N],
+            num_warps=num_warps,
+        )
+
+        mma_barrier = ttgl.allocate_shared_memory(ttgl.int64, [1], mbarrier.MBarrierLayout())
+        mbarrier.init(mma_barrier, count=1)
+
+        acc_zero = ttgl.zeros([M, N], dtype=acc_dtype, layout=tmem_reg_layout)
+        acc_tmem = allocate_tensor_memory(acc_dtype, [M, N], tmem_layout, acc_zero)
 
-    if ASYNC:
-        acc = hopper.warpgroup_mma_wait(num_outstanding=0, deps=[acc])
+        tcgen05_mma(smem_a, smem_b, acc_tmem, use_acc=False)
+        tcgen05_commit(mma_barrier)
+        mbarrier.wait(mma_barrier, phase=0)
+        mbarrier.invalidate(mma_barrier)
+        acc = acc_tmem.load(tmem_reg_layout)
+        acc = ttgl.convert_layout(acc, layout=mma_layout)
+    else:
+        acc = ttgl.zeros([M, N], dtype=acc_dtype, layout=mma_layout)
+        acc = hopper.warpgroup_mma(smem_a, smem_b, acc, is_async=ASYNC)
+
+        if ASYNC:
+            acc = hopper.warpgroup_mma_wait(num_outstanding=0, deps=[acc])
 
     ttgl.store(out + out_offs_m * N + out_offs_n, acc)
 
@@ -168,7 +192,7 @@ def test_warpgroup_mma(ASYNC):
     a = torch.randn((M, K), device="cuda", dtype=torch.float16)
     b = torch.randn((K, N), device="cuda", dtype=torch.float16)
     out = torch.zeros((M, N), device="cuda", dtype=torch.float16)
-    warpgroup_mma_kernel[(1, )](
+    mma_kernel[(1, )](
         a,
         b,
         out,
@@ -181,6 +205,7 @@ def test_warpgroup_mma(ASYNC):
         shared_layout_b,
         ttgl.float16,
         ASYNC,
+        False,
         num_warps=warps[0] * warps[1],
     )
 
@@ -189,19 +214,24 @@ def test_warpgroup_mma(ASYNC):
     torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-1)
 
 
-@pytest.mark.xfail(not is_hopper(), reason="Requires Hopper", run=False)
+@pytest.mark.xfail(not (is_hopper() or is_blackwell()), reason="Requires Hopper or Blackwell", run=False)
 @pytest.mark.parametrize("bitwidth, transpose_a, transpose_b, acc_dtype",
                          [(bitwidth, transpose_a, transpose_b, acc_dtype)
                           for bitwidth in [8, 16, 32]
                           for (transpose_a, transpose_b) in product([False, True], repeat=2)
                           for acc_dtype in [torch.float16, torch.float32]
                           if bitwidth == 16 or (acc_dtype == torch.float32 and not transpose_a and transpose_b)])
 @pytest.mark.parametrize("warps", ([8, 1], [4, 2], [4, 1]))
-# Swizzling 0 does not map to a valid memory descriptor lol
-@pytest.mark.parametrize("swizzling_a, swizzling_b", product([32, 64, 128], repeat=2))
+@pytest.mark.parametrize("swizzling_a, swizzling_b", product([0, 32, 64, 128], repeat=2))
 @pytest.mark.parametrize("shape_m, shape_n, shape_k", [(1, 1, 1), (2, 4, 1), (2, 2, 4)])
-def test_warpgroup_mma_shared_inputs(bitwidth, transpose_a, transpose_b, acc_dtype, warps, swizzling_a, swizzling_b,
-                                     shape_m, shape_n, shape_k):
+def test_mma_shared_inputs(bitwidth, transpose_a, transpose_b, acc_dtype, warps, swizzling_a, swizzling_b, shape_m,
+                           shape_n, shape_k, fresh_knobs):
+
+    # FIXME: Workaround for a bug in PTXAS when the shared layout is transposed and the swizzling is 0
+    if bitwidth == 16 and ((transpose_a and swizzling_a == 0 and shape_m > 1) or
+                           (not transpose_b and swizzling_b == 0 and shape_n > 1)):
+        fresh_knobs.nvidia.disable_ptxas_opt = True
+    use_tcgen05 = is_blackwell()
 
     torch_dtype_map = {
         8: torch.float8_e4m3fn,
@@ -214,8 +244,7 @@ def test_warpgroup_mma_shared_inputs(bitwidth, transpose_a, transpose_b, acc_dty
     }
 
     # We'll choose a larger instr shape along N, but sure
-    instr_shape_k_map = {8: 32, 16: 16, 32: 8}
-    instr_shape = [16, 32, instr_shape_k_map[bitwidth]]
+    instr_shape = [16, 32, 256 // bitwidth]
     M = instr_shape[0] * warps[0]
     N = instr_shape[1] * warps[1]
     K = instr_shape[2]
@@ -239,7 +268,27 @@ def min_shape(swizzling, dim0, dim1, trans):
     K *= shape_k
     instr_shape[1] *= shape_n
 
-    shared_mem_accum = M * K * bitwidth // 8 + K * N * bitwidth // 8
+    if use_tcgen05:
+        M = 128
+
+    def get_shared_swizzling_zero(M, K, transpose):
+        # K-contig
+        if transpose:
+            K, M = M, K
+        bases = []
+        for i in range(int(math.log2(128 // bitwidth))):
+            bases.append([0, 1 << i])
+        for i in range(int(math.log2(M))):
+            bases.append([1 << i, 0])
+        for i in range(int(math.log2(K // (128 // bitwidth)))):
+            offset = int(math.log2(128 // bitwidth)) + i
+            bases.append([0, 1 << offset])
+        if transpose:
+            for i in range(len(bases)):
+                bases[i] = [bases[i][1], bases[i][0]]
+        return ttgl.SharedLinearLayout(bases)
+
+    shared_mem_accum = (M + N) * K * bitwidth // 8
     if triton.runtime.driver.active.utils.get_device_properties(
             triton.runtime.driver.active.get_current_device())["max_shared_mem"] < shared_mem_accum:
         pytest.skip("Skipped due to insufficient shared memory on this GPU.")
@@ -248,11 +297,17 @@ def min_shape(swizzling, dim0, dim1, trans):
     gl_acc_dtype = acc_dtype_map[acc_dtype]
     out_dtype = torch.float32
 
-    block_layout = ttgl.BlockedLayout([1, 1], [1, THREADS_PER_WARP], warps_per_cta=warps, order=[1, 0])
-    shared_layout_a = ttgl.NVMMASharedLayout(swizzle_byte_width=swizzling_a, element_bitwidth=bitwidth, rank=2,
-                                             transposed=transpose_a)
-    shared_layout_b = ttgl.NVMMASharedLayout(swizzle_byte_width=swizzling_b, element_bitwidth=bitwidth, rank=2,
-                                             transposed=transpose_b)
+    block_layout = ttgl.BlockedLayout([1, 8], [1, THREADS_PER_WARP], warps_per_cta=warps, order=[1, 0])
+    if swizzling_a == 0:
+        shared_layout_a = get_shared_swizzling_zero(M, K, transpose_a)
+    else:
+        shared_layout_a = ttgl.NVMMASharedLayout(swizzle_byte_width=swizzling_a, element_bitwidth=bitwidth, rank=2,
+                                                 transposed=transpose_a)
+    if swizzling_b == 0:
+        shared_layout_b = get_shared_swizzling_zero(K, N, transpose_b)
+    else:
+        shared_layout_b = ttgl.NVMMASharedLayout(swizzle_byte_width=swizzling_b, element_bitwidth=bitwidth, rank=2,
+                                                 transposed=transpose_b)
     mma_layout = ttgl.NVMMADistributedLayout(version=[3, 0], warps_per_cta=warps, instr_shape=instr_shape)
 
     torch.manual_seed(0)
@@ -271,7 +326,7 @@ def cast(x, dtype):
     b = cast(torch.randn((K, N), device="cuda", dtype=torch.float32), torch_dtype)
     out = torch.zeros((M, N), device="cuda", dtype=out_dtype)
 
-    warpgroup_mma_kernel[(1, )](
+    mma_kernel[(1, )](
         a,
         b,
         out,
@@ -284,6 +339,7 @@ def cast(x, dtype):
         shared_layout_b,
         gl_acc_dtype,
         False,
+        use_tcgen05,
         num_warps=warps[0] * warps[1],
     )
 
@@ -298,9 +354,9 @@ def cast(x, dtype):
         torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = allow_fp16_red
 
     if bitwidth == 8:
-        atol, rtol = 0.5, 0.5
+        atol, rtol = 5e-2, 5e-1
     elif bitwidth == 16:
-        atol, rtol = 3e-2, 1e-1
+        atol, rtol = 5e-2, 5e-1
     else:
         atol, rtol = 5e-4, 5e-3
     torch.testing.assert_close(out, ref, atol=atol, rtol=rtol)