Merge commit 'cfe3dd079865d9f23cc6d7854de707636892cacf'

Author: whitneywhtsang

include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h

@@ -4,6 +4,7 @@
 #include "triton/Conversion/MLIRTypes.h"
 
 namespace mlir::triton {
+enum class ProgramIDDim : uint32_t;
 
 class TargetInfoBase {
 public:
@@ -48,7 +49,7 @@ class TargetInfoBase {
                            Value i) const = 0;
 
   virtual Value programId(RewriterBase &rewriter, Location loc,
-                          ModuleOp moduleOp, int axis) const = 0;
+                          ModuleOp moduleOp, ProgramIDDim axis) const = 0;
 
   virtual bool warpReduce(RewriterBase &rewriter, Location loc,
                           SmallVector<Value> &acc, triton::ReduceOp op,

lib/Conversion/TritonGPUToLLVM/PrintOpToLLVM.cpp

@@ -26,11 +26,10 @@ struct PrintOpConversion : public ConvertOpToLLVMPattern<triton::PrintOp> {
                   ConversionPatternRewriter &rewriter) const override {
     auto loc = op->getLoc();
 
-    auto getPid = [&](int axis) {
-      return targetInfo.programId(rewriter, loc,
-                                  op->getParentOfType<ModuleOp>(), axis);
-    };
-    std::array<Value, 3> pid = {getPid(0), getPid(1), getPid(2)};
+    std::array<Value, 3> pid;
+    auto module = op->getParentOfType<ModuleOp>();
+    for (auto axis : {ProgramIDDim::X, ProgramIDDim::Y, ProgramIDDim::Z})
+      pid[(int)axis] = targetInfo.programId(rewriter, loc, module, axis);
 
     // Simple printf of a string without any tensors.
     if (op.getNumOperands() == 0) {

lib/Conversion/TritonGPUToLLVM/SPMDOpToLLVM.cpp

@@ -17,9 +17,8 @@ struct GetProgramIdOpConversion
   LogicalResult
   matchAndRewrite(triton::GetProgramIdOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    Value programId = targetInfo.programId(rewriter, op->getLoc(),
-                                           op->getParentOfType<ModuleOp>(),
-                                           op.getAxisAsInt());
+    Value programId = targetInfo.programId(
+        rewriter, op->getLoc(), op->getParentOfType<ModuleOp>(), op.getAxis());
     rewriter.replaceOp(op, programId);
     return success();
   }

lib/Dialect/Gluon/Transforms/ResolveAutoEncodings.cpp

@@ -1,3 +1,4 @@
+#include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Visitors.h"
 #include "mlir/Support/LLVM.h"
@@ -9,6 +10,8 @@
 #include "llvm/ADT/PriorityWorklist.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/LogicalResult.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/Support/xxhash.h"
 
 namespace ttg = mlir::triton::gpu;
 
@@ -28,6 +31,65 @@ bool isAutoEncodingTensorType(Type ty) {
   return tensorTy && isa<gluon::AutoEncodingAttr>(tensorTy.getEncoding());
 }
 
+struct LayoutInfo {
+  Attribute encoding;
+  // Some operations can infer one of many encodings,
+  // we model this by setting the mayVary flag on encodings
+  // derived from these ops.
+  // If "may vary" is set then we allow conflicts, and when
+  // resolving conflicts we prefer encodings that are not allowed to vary.
+  bool mayVary = false;
+
+  operator bool() { return bool(encoding); }
+};
+
+uint64_t hashWithMemo(Attribute attr,
+                      llvm::MapVector<Attribute, uint64_t> &hashMemo) {
+  auto it = hashMemo.find(attr);
+  if (it != hashMemo.end()) {
+    return it->second;
+  }
+
+  // llvm::hash_value is not stable, so instead we hash the string repr of the
+  // attribute
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  attr.print(os);
+  auto hash = llvm::xxh3_64bits(str);
+  hashMemo.try_emplace(attr, hash);
+  return hash;
+}
+
+bool compare(Attribute a, Attribute b,
+             llvm::MapVector<Attribute, uint64_t> &hashMemo) {
+  if (a == b)
+    return false;
+
+  return hashWithMemo(a, hashMemo) > hashWithMemo(b, hashMemo);
+}
+
+LayoutInfo combineInfo(LayoutInfo lhs, LayoutInfo rhs, Operation *op,
+                       llvm::MapVector<Attribute, uint64_t> &hashMemo) {
+  // Sort inputs so this operation is commutative
+  if (compare(lhs.encoding, rhs.encoding, hashMemo)) {
+    std::swap(lhs, rhs);
+  }
+  if (lhs.mayVary)
+    return rhs;
+  if (rhs.mayVary)
+    return lhs;
+  if (lhs.encoding == rhs.encoding)
+    return lhs;
+  op->emitOpError("found conflicting encodings for value:\n  ")
+      << lhs.encoding << "\nand\n  " << rhs.encoding;
+  return {};
+}
+
+bool encodingsMayVary(Operation *op) {
+  return isa<triton::JoinOp, triton::SplitOp, triton::ReshapeOp, triton::CatOp,
+             triton::TransOp>(op);
+}
+
 LogicalResult inferAutoLayouts(FuncOp func) {
   // Disallow auto encoding accross function call boundaries
   for (auto argTy : func.getArgumentTypes()) {
@@ -42,33 +104,37 @@ LogicalResult inferAutoLayouts(FuncOp func) {
           "Functions returning auto encoding must be fully inlined");
   }
 
-  llvm::MapVector<Value, Attribute> valueToEncoding;
+  llvm::MapVector<Value, LayoutInfo> valueToEncoding;
   llvm::PriorityWorklist<Value> worklist;
+  llvm::MapVector<Attribute, uint64_t> hashMemo;
 
   auto updateEncoding = [&](ArrayRef<Value> values,
-                            Attribute enc) -> LogicalResult {
+                            LayoutInfo info) -> LogicalResult {
     for (auto value : values) {
-      auto [it, inserted] = valueToEncoding.insert({value, enc});
+      auto [it, inserted] = valueToEncoding.insert({value, info});
       if (!inserted) {
-        if (it->second != enc) {
-          auto defOp = value.getDefiningOp();
-          auto op = defOp ? defOp : func;
-          return op->emitOpError("Found conflicting encodings for value");
-        }
-      } else {
-        LLVM_DEBUG({
-          DBGS() << "Setting value:\n\t" << value << "\nto encoding:\n\t" << enc
-                 << "\n";
-        });
-        worklist.insert(value);
+        auto defOp = value.getDefiningOp();
+        auto op = defOp ? defOp : func;
+        auto combine = combineInfo(it->second, info, op, hashMemo);
+        if (!combine)
+          return failure();
+        if (combine == it->second)
+          continue;
+        it->second = combine;
       }
+      LLVM_DEBUG({
+        DBGS() << "Setting value:\n\t" << value << "\nto encoding:\n\t"
+               << it->second << "\n";
+      });
+      worklist.insert(value);
     }
     return success();
   };
 
   // 1. Set seed values from set_auto_layout ops
   auto res = func.walk([&](gluon::SetAutoLayoutOp op) -> WalkResult {
-    return updateEncoding({op.getSrc()}, op.getType().getEncoding());
+    return updateEncoding({op.getSrc()},
+                          LayoutInfo{op.getType().getEncoding()});
   });
 
   if (res.wasInterrupted())
@@ -77,26 +143,28 @@ LogicalResult inferAutoLayouts(FuncOp func) {
   // 2. Propagate encodings through the graph until fixed point, or conflict
   while (!worklist.empty()) {
     auto val = worklist.pop_back_val();
-    auto enc = valueToEncoding[val];
-    assert(enc);
+    auto info = valueToEncoding[val];
+    assert(info);
 
     // Propagate to users
     for (OpOperand &use : val.getUses()) {
       auto op = use.getOwner();
       if (isa<scf::ForOp, scf::WhileOp>(op)) {
         auto offset = 3 * isa<scf::ForOp>(op);
         auto tiedArgs = getTiedArgs(op, use.getOperandNumber() - offset);
-        if (failed(updateEncoding(tiedArgs, enc)))
+        if (failed(updateEncoding(tiedArgs, info)))
           return failure();
       } else if (isa<scf::YieldOp>(op)) {
         auto tiedArgs = getTiedArgs(op, use.getOperandNumber());
-        if (failed(updateEncoding(tiedArgs, enc)))
+        if (failed(updateEncoding(tiedArgs, info)))
           return failure();
       } else {
-        auto dstEnc = inferDstEncoding(op, enc);
+        auto dstEnc = inferDstEncoding(op, info.encoding);
         if (dstEnc) {
+          bool mayVary = info.mayVary || encodingsMayVary(op);
+          LayoutInfo dstInfo{dstEnc, mayVary};
           if (failed(updateEncoding(llvm::to_vector_of<Value>(op->getResults()),
-                                    dstEnc)))
+                                    dstInfo)))
             return failure();
         }
       }
@@ -107,17 +175,19 @@ LogicalResult inferAutoLayouts(FuncOp func) {
       auto definingOp = opResult.getOwner();
       if (isa<scf::ForOp, scf::WhileOp, scf::IfOp>(definingOp)) {
         auto tiedArgs = getTiedArgs(definingOp, opResult.getResultNumber());
-        if (failed(updateEncoding(tiedArgs, enc)))
+        if (failed(updateEncoding(tiedArgs, info)))
           return failure();
       } else {
-        auto srcEncoding = inferSrcEncoding(definingOp, enc);
+        auto srcEncoding = inferSrcEncoding(definingOp, info.encoding);
         if (srcEncoding) {
+          bool mayVary = info.mayVary || encodingsMayVary(definingOp);
+          LayoutInfo srcInfo{srcEncoding, mayVary};
           llvm::SmallVector<Value> tensorOperands;
           for (auto operand : definingOp->getOperands())
             if (isa<RankedTensorType>(operand.getType()))
               tensorOperands.push_back(operand);
 
-          if (failed(updateEncoding(tensorOperands, srcEncoding)))
+          if (failed(updateEncoding(tensorOperands, srcInfo)))
             return failure();
         }
       }
@@ -126,18 +196,18 @@ LogicalResult inferAutoLayouts(FuncOp func) {
       if (isa<scf::ForOp, scf::WhileOp>(parentOp)) {
         auto offset = isa<scf::ForOp>(parentOp);
         auto tiedArgs = getTiedArgs(parentOp, blockArg.getArgNumber() - offset);
-        if (failed(updateEncoding(tiedArgs, enc)))
+        if (failed(updateEncoding(tiedArgs, info)))
           return failure();
       }
     }
   }
 
   // 3. Transfer propagated encodings into the graph
   auto ctx = func.getContext();
-  for (auto &[val, enc] : valueToEncoding) {
+  for (auto &[val, info] : valueToEncoding) {
     auto existingTy = cast<RankedTensorType>(val.getType());
     assert(isa<gluon::AutoEncodingAttr>(existingTy.getEncoding()));
-    auto ty = existingTy.cloneWithEncoding(enc);
+    auto ty = existingTy.cloneWithEncoding(info.encoding);
     val.setType(ty);
 
     if (auto opResult = dyn_cast<OpResult>(val)) {

python/test/unit/language/test_core.py

@@ -4527,6 +4527,9 @@ def make_finite(x, dtype):
             assert 'st.global.v4' in ptx
         assert (re.search(r'(mma|wgmma.mma_async).sync.aligned.m\d+n\d+k16(?:.row.col)?.f32.(f|bf)16.(f|bf)16', ptx)
                 or "tcgen05.mma.cta_group::1.kind::f16" in ptx)
+    if is_hip_cdna4() and normal_type in ["bf16", "fp16"]:
+        amdgcn = pgm.asm['amdgcn']
+        assert (re.search(r"v_cvt_scalef32_pk_.*?(fp4|fp8|bf8).*?op_sel", amdgcn))
 
 
 @pytest.mark.interpreter

python/triton/language/__init__.py

@@ -84,7 +84,7 @@
     join,
     load,
     make_block_ptr,
-    map_elementwise,  # noqa
+    map_elementwise,
     max_constancy,
     max_contiguous,
     maximum,
@@ -209,6 +209,7 @@
     "log",
     "log2",
     "make_block_ptr",
+    "map_elementwise",
     "math",
     "max",
     "max_constancy",

python/triton/runtime/interpreter.py

@@ -93,13 +93,6 @@ def validate(self):
         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/Analysis/amd/test-alignment.mlir

@@ -2,7 +2,7 @@
 
 #mma = #ttg.amd_mfma<{version = 3, warpsPerCTA = [4, 1], instrShape = [32, 32], isTransposed = true}>
 
-tt.func public @kernel(%arg0: tensor<256x64xf16, #mma> {tt.contiguity=256 : i32, tt.divisibility=6: i32, tt.constancy=1: i32}) attributes {noinline = false} {
+tt.func public @kernel(%arg0: tensor<256x64xf16, #mma> {tt.contiguity=256 : i32, tt.divisibility=6: i32, tt.constancy=1: i32}) {
   // expeted-remark @below {{contiguity = [128, 32], divisibility = [6, 6], constancy = [1, 1], constant_value = <none>}}
   %0 = amdgpu.extract_slice %arg0 [128, 32] : tensor<256x64xf16, #mma> to tensor<128x32xf16, #mma>
   tt.return

test/Conversion/amd/buffer_atomic_cas.mlir

@@ -2,7 +2,7 @@
 #blocked = #ttg.blocked<{sizePerThread = [2], threadsPerWarp = [64], warpsPerCTA = [4], order = [0]}>
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "hip:gfx942", "ttg.threads-per-warp" = 64 : i32} {
   // CHECK-LABEL: buffer_atomic_cas_i64
-  tt.func public @buffer_atomic_cas_i64(%arg0: !tt.ptr<i64> {tt.divisibility = 16 : i32, tt.pointer_range = 32 : i32}, %arg1: !tt.ptr<i64> {tt.divisibility = 16 : i32, tt.pointer_range = 32 : i32}) attributes {noinline = false} {
+  tt.func public @buffer_atomic_cas_i64(%arg0: !tt.ptr<i64> {tt.divisibility = 16 : i32, tt.pointer_range = 32 : i32}, %arg1: !tt.ptr<i64> {tt.divisibility = 16 : i32, tt.pointer_range = 32 : i32}) {
     // CHECK: %[[cas_val:.*]] = llvm.mlir.constant(2 : i64) : i64
     // CHECK: %[[cas_val_cast:.*]] = llvm.bitcast %[[cas_val]] : i64 to i64
     // CHECK: %[[cas_val_insert:.*]] = llvm.insertvalue %[[cas_val_cast]], %{{.*}}[1] : !llvm.struct<(i64, i64)>

test/Conversion/amd/buffer_load_store.mlir

@@ -262,7 +262,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.thr
 #blocked = #ttg.blocked<{sizePerThread = [4], threadsPerWarp = [64], warpsPerCTA = [4], order = [0]}>
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 64 : i32} {
   // CHECK-LABEL: strided_buffer_load_and_store
-  tt.func public @strided_buffer_load_and_store(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32, tt.pointer_range = 32 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32, tt.pointer_range = 32 : i32}) attributes {noinline = false} {
+  tt.func public @strided_buffer_load_and_store(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32, tt.pointer_range = 32 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32, tt.pointer_range = 32 : i32}) {
     %cst = arith.constant dense<2> : tensor<1024xi32, #blocked>
     %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32, #blocked>
     %1 = arith.muli %0, %cst : tensor<1024xi32, #blocked>