Merge commit '4bcdbdee14a992b6086afb1f3025fe2767fdbbec'

Author: whitneywhtsang

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

@@ -87,7 +87,7 @@ def TTNG_WarpGroupDotOp : TTNG_Op<"warp_group_dot", [
 
   let arguments = (ins
     TTG_TensorOrMemDesc:$a,
-    TTG_TensorOrMemDesc:$b,
+    TTG_MemDescType:$b,
     TT_FpIntTensor:$c,
     Optional<I1>:$useC,
     DefaultValuedAttr<TT_InputPrecisionAttr, "::mlir::triton::InputPrecision::IEEE">:$inputPrecision,
@@ -99,7 +99,7 @@ def TTNG_WarpGroupDotOp : TTNG_Op<"warp_group_dot", [
 
   let assemblyFormat = [{
     $a`,` $b`,` $c (`,` $useC^)? attr-dict
-    `:` type($a) `*` type($b) `->` type($d)
+    `:` type($a) `*` qualified(type($b)) `->` type($d)
   }];
 
   let extraClassDeclaration = [{

lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp

@@ -63,9 +63,17 @@ LogicalResult WarpGroupDotOp::verify() {
   auto nvmmaEnc = dyn_cast<NvidiaMmaEncodingAttr>(resTy.getEncoding());
   if (!nvmmaEnc || !nvmmaEnc.isHopper())
     return emitOpError("WGMMA result layout must be Hopper NVMMA");
+
+  if (!isa<NVMMASharedEncodingAttr, DotOperandEncodingAttr>(
+          getA().getType().getEncoding()))
+    return emitOpError("WGMMA A operand must have NVMMA shared or dot layout");
+  if (!isa<NVMMASharedEncodingAttr>(getB().getType().getEncoding()))
+    return emitOpError("WGMMA B operand must have NVMMA shared layout");
+
   auto numWarps = gpu::lookupNumWarps(getOperation());
   if (numWarps % 4)
     return emitOpError("WGMMA requires num_warps to be divisible by 4");
+
   auto retShapePerCTA = getShapePerCTA(resTy);
   int rank = retShapePerCTA.size();
   if (rank != 2)
@@ -74,12 +82,14 @@ LogicalResult WarpGroupDotOp::verify() {
     return emitOpError("WGMMA result M dimension must be divisible by 64");
   if (retShapePerCTA[1] % 8 != 0)
     return emitOpError("WGMMA result N dimension must be divisible by 8");
+
   auto aElemTy = getA().getType().getElementType();
   if (!(llvm::isa<Float8E5M2Type, Float8E4M3FNType>(aElemTy) ||
         aElemTy.isInteger(8) || aElemTy.isF16() || aElemTy.isBF16() ||
         aElemTy.isF32()))
     return emitOpError("WGMMA result element type must be F16, BF16, F32, "
                        "F8E5M2, F8E4M3FN, or integer type");
+
   if (getMaxNumImpreciseAcc() < 32 &&
       (llvm::isa<Float8E5M2Type, Float8E4M3FNType>(aElemTy)) &&
       resTy.getElementType().isF32()) {

lib/Dialect/TritonNvidiaGPU/Transforms/TensorMemoryAllocation.cpp

@@ -1,9 +1,13 @@
 #include "mlir/Analysis/Liveness.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 #include "triton/Analysis/Allocation.h"
 #include "triton/Dialect/Triton/IR/Utility.h"
+#include "triton/Dialect/TritonGPU/IR/Dialect.h"
+#include "triton/Dialect/TritonGPU/IR/Traits.h"
 #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h"
 #include "llvm/ADT/EquivalenceClasses.h"
@@ -175,30 +179,92 @@ static TMemChunk allocFirstFit(MemoryBitMap &memoryMap,
   return chunk;
 }
 
-static Operation *getAlloc(Value value) {
-  while (true) {
-    if (auto allocOp = value.getDefiningOp<TMEMAllocOp>())
-      return allocOp;
-    if (auto indexOp = value.getDefiningOp<ttg::MemDescIndexOp>()) {
-      value = indexOp.getSrc();
+static SmallVector<Operation *> getAlloc(Value value) {
+  SmallVector<Operation *> allocs;
+  DenseSet<Value> seen;
+  SmallVector<Value> worklist{value};
+
+  while (!worklist.empty()) {
+    Value v = worklist.pop_back_val();
+    if (!seen.insert(v).second)
       continue;
-    }
-    if (auto reinterpOp = value.getDefiningOp<ttg::MemDescReinterpretOp>()) {
-      value = reinterpOp.getSrc();
+
+    // Handle block arguments.
+    if (auto arg = dyn_cast<BlockArgument>(v)) {
+      Block *block = arg.getOwner();
+      Operation *parentOp = block->getParentOp();
+
+      // Handle block with predecessors.
+      if (!block->isEntryBlock()) {
+        for (Block *pred : block->getPredecessors()) {
+          Operation *predOp = pred->getTerminator();
+          auto br = dyn_cast<BranchOpInterface>(predOp);
+          if (!br) {
+            llvm::report_fatal_error("unhandled branch op: " +
+                                     predOp->getName().getStringRef());
+          }
+          SmallVector<Attribute> operands(br->getNumOperands());
+          auto it = llvm::find(br->getSuccessors(), block);
+          unsigned idx = std::distance(br->getSuccessors().begin(), it);
+          SuccessorOperands args = br.getSuccessorOperands(idx);
+          Value operand =
+              args.getForwardedOperands()[arg.getArgNumber() -
+                                          args.getProducedOperandCount()];
+          worklist.push_back(operand);
+        }
+        continue;
+      }
+
+      // Handle region entry arguments.
+      if (auto wsOp = dyn_cast<ttg::WarpSpecializePartitionsOp>(parentOp)) {
+        worklist.push_back(
+            wsOp.getParentOp().getExplicitCaptures()[arg.getArgNumber()]);
+      } else if (auto forOp = dyn_cast<scf::ForOp>(parentOp)) {
+        unsigned idx = arg.getArgNumber() - 1;
+        worklist.push_back(forOp.getYieldedValues()[idx]);
+        worklist.push_back(forOp.getInits()[idx]);
+      } else if (auto whileOp = dyn_cast<scf::WhileOp>(parentOp)) {
+        unsigned idx = arg.getArgNumber();
+        if (arg.getParentRegion() == &whileOp.getAfter()) {
+          worklist.push_back(whileOp.getConditionOp().getArgs()[idx]);
+        } else {
+          worklist.push_back(whileOp.getYieldedValues()[idx]);
+          worklist.push_back(whileOp.getInits()[idx]);
+        }
+      } else {
+        llvm::report_fatal_error(
+            "unhandled parent op when looking for TMEM alloc: " +
+            parentOp->getName().getStringRef());
+      }
       continue;
     }
-    if (auto slice = value.getDefiningOp<TMEMSubSliceOp>()) {
-      value = slice.getSrc();
-      continue;
+
+    Operation *defOp = v.getDefiningOp();
+    unsigned idx = cast<OpResult>(v).getResultNumber();
+    if (isa<TMEMAllocOp>(defOp)) {
+      allocs.push_back(defOp);
+    } else if (defOp->hasTrait<OpTrait::MemDescViewTrait>()) {
+      worklist.push_back(defOp->getOperand(0));
+    } else if (auto sliceOp = dyn_cast<TMEMSubSliceOp>(defOp)) {
+      worklist.push_back(sliceOp.getSrc());
+    } else if (auto selectOp = dyn_cast<arith::SelectOp>(defOp)) {
+      worklist.push_back(selectOp.getTrueValue());
+      worklist.push_back(selectOp.getFalseValue());
+    } else if (auto ifOp = dyn_cast<scf::IfOp>(defOp)) {
+      worklist.push_back(ifOp.thenYield().getOperand(idx));
+      worklist.push_back(ifOp.elseYield().getOperand(idx));
+    } else if (auto forOp = dyn_cast<scf::ForOp>(defOp)) {
+      worklist.push_back(forOp.getYieldedValues()[idx]);
+      worklist.push_back(forOp.getInits()[idx]);
+    } else if (auto whileOp = dyn_cast<scf::WhileOp>(defOp)) {
+      worklist.push_back(whileOp.getConditionOp().getArgs()[idx]);
+    } else {
+      llvm::report_fatal_error("unhandled op when looking for TMEM alloc: " +
+                               defOp->getName().getStringRef());
     }
-    auto arg = dyn_cast<BlockArgument>(value);
-    if (!arg || !isa<triton::gpu::WarpSpecializePartitionsOp>(
-                    arg.getOwner()->getParentOp()))
-      llvm::report_fatal_error("expected to find a TMEM alloc op");
-    auto partitions = cast<triton::gpu::WarpSpecializePartitionsOp>(
-        arg.getOwner()->getParentOp());
-    value = partitions.getParentOp().getExplicitCaptures()[arg.getArgNumber()];
   }
+
+  return allocs;
 }
 
 class RowIdConstraints {
@@ -245,8 +311,11 @@ allocateTMem(Operation *parentOp,
         if (allocSize.numRows == 64) {
           // HW restriction, the A alloc and accumulator needs to be in the same
           // rows.
-          rowIdConstraints.joinOps(getAlloc(mmaOp.getA()),
-                                   getAlloc(mmaOp.getAccumulator()));
+          SmallVector<Operation *> lhsAllocs = getAlloc(mmaOp.getA());
+          SmallVector<Operation *> accAllocs = getAlloc(mmaOp.getAccumulator());
+          for (Operation *lhsAlloc : lhsAllocs)
+            for (Operation *accAlloc : accAllocs)
+              rowIdConstraints.joinOps(lhsAlloc, accAlloc);
         } else {
           // TODO: we need to handle cases where the format is blockM and we
           // have multiple blocks.

python/test/unit/language/test_core.py

@@ -1969,24 +1969,33 @@ def serialized_add(data, Lock, SEM: tl.constexpr):
 
 
 @pytest.mark.interpreter
-@pytest.mark.parametrize("sem", [None, 'acquire', 'release', 'acq_rel', 'relaxed'])
+@pytest.mark.parametrize("sem", [None, "acquire", "release", "acq_rel", "relaxed"])
 @pytest.mark.parametrize("num_ctas", num_ctas_list)
-def test_tensor_atomic_cas(sem, num_ctas, device):
+@pytest.mark.parametrize("size", [4, 128, 512])
+@pytest.mark.parametrize("dtype_str", ['bfloat16', 'float16', 'float32', 'uint64', 'int64', 'float64'])
+def test_tensor_atomic_cas(sem, size, dtype_str, num_ctas, device):
+    check_type_supported(dtype_str, device)
+    if "float" in dtype_str and is_hip():
+        pytest.skip("HIP does not support atomic cas with float types")
 
     @triton.jit
-    def change_value(X, BLOCK_SIZE: tl.constexpr, sem: tl.constexpr):
+    def change_value(X, BLOCK_SIZE: tl.constexpr, sem: tl.constexpr, dtype: tl.constexpr):
         pid = tl.program_id(axis=0)
         block_start = pid * BLOCK_SIZE
         offsets = block_start + tl.arange(0, BLOCK_SIZE)
-        t1 = tl.full((BLOCK_SIZE, ), 0, dtype=tl.int64)
-        t2 = tl.full((BLOCK_SIZE, ), 2, dtype=tl.int64)
+        t1 = tl.full((BLOCK_SIZE, ), 0, dtype=dtype)
+        t2 = tl.full((BLOCK_SIZE, ), 2, dtype=dtype)
         tl.atomic_cas(X + offsets, t1, t2, sem=sem)
 
-    X = torch.tensor([0, 1, 0, 1, 0, 1, 0, 1], device=device, dtype=torch.int64)
-    Y = torch.tensor([2, 1, 2, 1, 2, 1, 2, 1], device=device, dtype=torch.int64)
+    torch_dtype = getattr(torch, dtype_str)
+    X = torch.zeros((size, ), device=device, dtype=torch_dtype)
+    X[1::2] = 1
+    Y = X.clone()
+    Y[0::2] = 2
 
-    change_value[(2, )](X, 4, sem)
-    assert (torch.equal(X, Y))
+    tl_dtype = getattr(tl, dtype_str)
+    change_value[(2, )](X, BLOCK_SIZE=size // 2, sem=sem, dtype=tl_dtype)
+    assert torch.equal(X, Y)
 
 
 @pytest.mark.interpreter

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

@@ -26,7 +26,7 @@ def warpgroup_mma(a, b, acc, *, use_acc=True, precision=None, max_num_imprecise_
 
     Args:
         a (tensor or shared_memory_descriptor): Left hand side operand.
-        b (tensor or shared_memory_descriptor): Right hand side operand.
+        b (shared_memory_descriptor): Right hand side operand.
         acc (tensor): Accumulator tensor.
         use_acc (bool): Whether to use the initial value of the accumulator. Defaults to True.
         precision (str, optional): Dot input precision. Defaults to builder default.

python/triton/knobs.py

@@ -518,7 +518,6 @@ class amd_knobs(base_knobs):
     use_buffer_atomics: env_bool = env_bool("AMDGCN_USE_BUFFER_ATOMICS", True)
     dump_amdgcn: env_bool = env_bool("AMDGCN_ENABLE_DUMP")
     libhip_path: env_opt_str = env_opt_str("TRITON_LIBHIP_PATH")
-    lld_path: env_opt_str = env_opt_str("TRITON_HIP_LLD_PATH")
 
     # We use strs so that we can have a default value based on other runtime info
     use_block_pingpong: env_opt_bool = env_opt_bool("TRITON_HIP_USE_BLOCK_PINGPONG")

python/triton/runtime/interpreter.py

@@ -88,10 +88,18 @@ def validate(self):
         assert self.base.data.item() % 16 == 0, "base must be 16-byte aligned"
         assert len(self.strides) == self.ndim
         assert len(self.block_shape) == self.ndim
+        assert self.ndim >= 1, "descriptor cannot be 0 dimensional"
 
         for stride in self.strides[:-1]:
             assert stride.data.item() % 16 == 0, "stride must be 16-byte aligned"
         assert self.strides[-1].data.item() == 1, "last dim must be contiguous"
+        for i in range(self.ndim - 1):
+            stride = self.strides[i].data.item()
+            prev_stride = self.strides[i + 1].data.item()
+            prev_size = self.shape[i + 1].data.item()
+            assert stride >= prev_stride, "strides must be ordered largest to smallest"
+            assert (stride % prev_stride) == 0, "strides must be even multiples of smaller strides"
+            assert (stride // prev_stride) >= prev_size, "invalid stride"
 
     def materialize_pointers(self, offsets: List[TensorHandle]):
         assert len(offsets) == self.ndim

test/TritonGPU/amd/accelerate-amd-matmul-mfma-gfx950.mlir

@@ -223,3 +223,74 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.targ
     tt.return
   }
 }
+
+// -----
+
+#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 64], warpsPerCTA = [2, 2], order = [1, 0]}>
+#blocked1 = #ttg.blocked<{sizePerThread = [16, 1], threadsPerWarp = [8, 8], warpsPerCTA = [1, 4], order = [0, 1]}>
+// CHECK{LITERAL}: #shared = #ttg.swizzled_shared<{vec = 16, perPhase = 4, maxPhase = 4, order = [1, 0]}>
+// CHECK-LABEL: mfma_dot_scaled_mxfp4_b_packed_mn
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "hip:gfx950", "ttg.threads-per-warp" = 64 : i32} {
+  tt.func public @mfma_dot_scaled_mxfp4_b_packed_mn(
+      %a: tensor<128x128xf8E5M2, #blocked>,
+      %b: tensor<128x64xi8, #blocked1>,
+      %c: tensor<128x128xf32, #blocked>,
+      %arg4: tensor<128x128x!tt.ptr<f32>, #blocked>
+      ) {
+    %b1 = ttg.convert_layout %b : tensor<128x64xi8, #blocked1> -> tensor<128x64xi8, #blocked>
+    // CHECK: %[[ALLOCB:.+]] = ttg.local_alloc {{.*}} : (tensor<128x64xi8, #blocked>) -> !ttg.memdesc<128x64xi8, #shared, #smem>
+    // CHECK: %[[B:.+]] = amdgpu.local_load_packed_tranposed  %[[ALLOCB]] : !ttg.memdesc<128x64xi8, #shared, #smem> -> tensor<64x128xi8, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 16}>>
+    // CHECK: tt.dot_scaled %{{.*}}, %[[B]], %{{.*}} lhs = e5m2 rhs = e2m1 {fastMath = false}
+    %accumulator_52 = tt.dot_scaled %a, %b1, %c lhs = e5m2 rhs = e2m1 {fastMath = false, rhs_k_pack = false} : tensor<128x128xf8E5M2, #blocked> * tensor<128x64xi8, #blocked> -> tensor<128x128xf32, #blocked>
+    tt.store %arg4, %accumulator_52 : tensor<128x128x!tt.ptr<f32>, #blocked>
+    tt.return
+  }
+}
+// -----
+
+#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 64], warpsPerCTA = [2, 2], order = [1, 0]}>
+#blocked1 = #ttg.blocked<{sizePerThread = [16, 1], threadsPerWarp = [8, 8], warpsPerCTA = [1, 4], order = [0, 1]}>
+// CHECK{LITERAL}: #shared = #ttg.swizzled_shared<{vec = 16, perPhase = 4, maxPhase = 4, order = [0, 1]}>
+// CHECK-LABEL: mfma_dot_scaled_mxfp4_a_packed_mn
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "hip:gfx950", "ttg.threads-per-warp" = 64 : i32} {
+  tt.func public @mfma_dot_scaled_mxfp4_a_packed_mn(
+      %a: tensor<64x128xi8, #blocked>,
+      %b: tensor<128x128xf8E5M2, #blocked1>,
+      %c: tensor<128x128xf32, #blocked>,
+      %arg4: tensor<128x128x!tt.ptr<f32>, #blocked>
+      ) {
+    %b1 = ttg.convert_layout %b : tensor<128x128xf8E5M2, #blocked1> -> tensor<128x128xf8E5M2, #blocked>
+    // CHECK: %[[ALLOCA:.+]] = ttg.local_alloc {{.*}} : (tensor<64x128xi8, #blocked>) -> !ttg.memdesc<64x128xi8, #shared, #smem>
+    // CHECK: %[[A:.+]] = amdgpu.local_load_packed_tranposed  %[[ALLOCA]] : !ttg.memdesc<64x128xi8, #shared, #smem> -> tensor<128x64xi8, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 16}>>
+    // CHECK: tt.dot_scaled %[[A]], %{{.*}}, %{{.*}} lhs = e2m1 rhs = e5m2 {fastMath = false}
+    %accumulator_52 = tt.dot_scaled %a, %b1, %c lhs = e2m1 rhs = e5m2 {fastMath = false, lhs_k_pack = false} : tensor<64x128xi8, #blocked> * tensor<128x128xf8E5M2, #blocked> -> tensor<128x128xf32, #blocked>
+    tt.store %arg4, %accumulator_52 : tensor<128x128x!tt.ptr<f32>, #blocked>
+    tt.return
+  }
+}
+
+// -----
+
+#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 64], warpsPerCTA = [2, 2], order = [1, 0]}>
+#blocked1 = #ttg.blocked<{sizePerThread = [16, 1], threadsPerWarp = [8, 8], warpsPerCTA = [1, 4], order = [0, 1]}>
+// CHECK{LITERAL}: #shared = #ttg.swizzled_shared<{vec = 16, perPhase = 4, maxPhase = 4, order = [0, 1]}>
+// CHECK{LITERAL}: #shared1 = #ttg.swizzled_shared<{vec = 16, perPhase = 4, maxPhase = 4, order = [1, 0]}>
+// CHECK-LABEL: mfma_dot_scaled_mxfp4_ab_packed_mn
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "hip:gfx950", "ttg.threads-per-warp" = 64 : i32} {
+  tt.func public @mfma_dot_scaled_mxfp4_ab_packed_mn(
+      %a: tensor<64x128xi8, #blocked>,
+      %b: tensor<128x64xi8, #blocked1>,
+      %c: tensor<128x128xf32, #blocked>,
+      %arg4: tensor<128x128x!tt.ptr<f32>, #blocked>
+      ) {
+    %b1 = ttg.convert_layout %b : tensor<128x64xi8, #blocked1> -> tensor<128x64xi8, #blocked>
+    // CHECK: %[[ALLOCA:.+]] = ttg.local_alloc {{.*}} : (tensor<64x128xi8, #blocked>) -> !ttg.memdesc<64x128xi8, #shared, #smem>
+    // CHECK: %[[A:.+]] = amdgpu.local_load_packed_tranposed  %[[ALLOCA]] : !ttg.memdesc<64x128xi8, #shared, #smem> -> tensor<128x64xi8, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 16}>>
+    // CHECK: %[[ALLOCB:.+]] = ttg.local_alloc {{.*}} : (tensor<128x64xi8, #blocked>) -> !ttg.memdesc<128x64xi8, #shared1, #smem>
+    // CHECK: %[[B:.+]] = amdgpu.local_load_packed_tranposed  %[[ALLOCB]] : !ttg.memdesc<128x64xi8, #shared1, #smem> -> tensor<64x128xi8, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 16}>>
+    // CHECK: tt.dot_scaled %[[A]], %[[B]], %{{.*}} lhs = e2m1 rhs = e2m1 {fastMath = false}
+    %accumulator_52 = tt.dot_scaled %a, %b1, %c lhs = e2m1 rhs = e2m1 {fastMath = false, lhs_k_pack = false, rhs_k_pack = false} : tensor<64x128xi8, #blocked> * tensor<128x64xi8, #blocked> -> tensor<128x128xf32, #blocked>
+    tt.store %arg4, %accumulator_52 : tensor<128x128x!tt.ptr<f32>, #blocked>
+    tt.return
+  }
+}