Merge OpenAI Triton commit 31baa6d (#5048)

This PR change the Triton base from
d412906
to 31baa6d (Sep 1).
Pass rate: 98.74%

---------

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

.github/workflows/build-macos.yml

@@ -106,6 +106,13 @@ jobs:
           source ~/.venv/bin/activate
           echo "PATH is '$PATH'"
           ccache --zero-stats
+          export PATH="/opt/homebrew/opt/llvm@19/bin:$PATH"
+          export CC="/opt/homebrew/opt/llvm@19/bin/clang"
+          export CXX="/opt/homebrew/opt/llvm@19/bin/clang++"
+          export CXXFLAGS="-stdlib=libc++"
+          export LDFLAGS="-L/opt/homebrew/opt/llvm@19/lib"
+          which clang++
+          clang++ --version
           make dev-install
       - name: CCache Stats
         run: ccache --print-stats

include/triton/Dialect/Triton/IR/TritonOpInterfaces.td

@@ -49,7 +49,12 @@ def DotOpInterface : OpInterface<"DotOpInterface"> {
       /*retType=*/"::mlir::Value",
       /*methodName=*/"getB",
       /*args=*/(ins)>,
-  InterfaceMethod<
+    InterfaceMethod<
+      /*desc=*/"Get the output tensor",
+      /*retType=*/"::mlir::Value",
+      /*methodName=*/"getD",
+      /*args=*/(ins)>,
+    InterfaceMethod<
       /*desc=*/"Verify the dimensions of the A and B DotOp operands.",
       /*retType=*/"bool",
       /*methodName=*/"verifyDims",
@@ -64,6 +69,7 @@ def DotOpInterface : OpInterface<"DotOpInterface"> {
         auto aTy = cast<ShapedType>($_op.getA().getType());
         auto bTy = cast<ShapedType>($_op.getB().getType());
         auto cTy = cast<ShapedType>($_op->getOperand(2).getType());
+        auto dTy = cast<ShapedType>($_op.getD().getType());
         auto aShape = aTy.getShape();
         auto bShape = bTy.getShape();
         auto cShape = cTy.getShape();

include/triton/Dialect/Triton/IR/TritonOps.td

@@ -1051,15 +1051,17 @@ def TT_MakeTensorDescOp : TT_Op<"make_tensor_descriptor", [
   let arguments = (ins
     TT_Ptr:$base,
     Variadic<I32>:$shape,
-    Variadic<I64>:$strides
+    Variadic<I64>:$strides,
+    DefaultValuedAttr<TT_PaddingOptionAttr, "::mlir::triton::PaddingOption::PAD_ZERO">:$padding
   );
 
   let results = (outs TT_TensorDescType:$result);
 
   let assemblyFormat = "$base `,` `[` $shape `]` `,` `[` $strides `]` attr-dict `:` type($base) `,` type($result)";
 
   let builders = [
-    OpBuilder<(ins "Value":$base, "ValueRange":$shape, "ValueRange":$strides, "ArrayRef<int32_t>":$blockShape, "bool":$isSignedInteger)>
+    OpBuilder<(ins "Value":$base, "ValueRange":$shape, "ValueRange":$strides, "ArrayRef<int32_t>":$blockShape, "bool":$isSignedInteger,
+    "triton::PaddingOption":$padding)>
   ];
 
   let extraClassDeclaration = [{

include/triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h

@@ -19,6 +19,7 @@ class AMDRotatingSharedEncodingAttr;
 class AMDMfmaEncodingAttr;
 class TensorOrMemDesc;
 class MemDescType;
+class CTALayoutAttr;
 
 // - BlockedEncodingAttrs have the following input dimensions.
 //
@@ -126,6 +127,13 @@ LinearLayout chooseScaledMfmaScaleLayout(MLIRContext *ctx, int dotOperandIdx,
                                          ArrayRef<unsigned> tilesPerWarp,
                                          ArrayRef<unsigned> warpsPerCTA);
 
+LinearLayout getSM120DotScaledScaleLayout(MLIRContext *ctx, int dotOperandIdx,
+                                          ArrayRef<int64_t> dotOperandShape,
+                                          ArrayRef<unsigned> tilesPerWarp,
+                                          ArrayRef<unsigned> warpsPerCTA,
+                                          unsigned instrM, unsigned instrN,
+                                          CTALayoutAttr ctaLayoutAttr);
+
 // Create LinearLayout for nvidia mma tile.
 LinearLayout nvidiaMmaTile(MLIRContext *ctx, ArrayRef<unsigned> tileShape,
                            unsigned kWidth, ArrayRef<unsigned> order,

lib/Analysis/Utility.cpp

@@ -649,6 +649,8 @@ bool supportMMA(triton::DotOp op, int version) {
   if (version == 5) {
     if (triton::tools::getBoolEnv("DISABLE_MMA_V5"))
       return false;
+    RankedTensorType typeA = op.getA().getType();
+    int k = typeA.getShape().back();
     auto retType = op.getType();
     auto retShapePerCTA = getShapePerCTA(retType);
     auto rank = retShapePerCTA.size();
@@ -662,8 +664,11 @@ bool supportMMA(triton::DotOp op, int version) {
       // Currently only support numWarps 4 or 8 for TMEM load and store.
       return false;
     }
+    // If k size is smaller than the native mma size, we cannot use MMA.
+    if (k < 256 / aElemTy.getIntOrFloatBitWidth())
+      return false;
     if (!(retShapePerCTA[rank - 2] % 64 == 0 &&
-          retShapePerCTA[rank - 1] % 8 == 0))
+          retShapePerCTA[rank - 1] % 16 == 0))
       return false;
     return true;
   }
@@ -683,7 +688,7 @@ bool supportMMA(triton::DotOp op, int version) {
     if (rank == 3)
       return false;
     if (!(numWarps % 4 == 0 && retShapePerCTA[rank - 2] % 64 == 0 &&
-          retShapePerCTA[rank - 1] % 8 == 0 &&
+          retShapePerCTA[rank - 1] % 16 == 0 &&
           (llvm::isa<Float8E5M2Type, Float8E4M3FNType>(aElemTy) ||
            aElemTy.isInteger(8) || aElemTy.isF16() || aElemTy.isBF16() ||
            aElemTy.isF32()))) {

lib/Conversion/TritonGPUToLLVM/HistogramOpToLLVM.cpp

@@ -25,7 +25,6 @@ static SmallVector<Value> computeWarpLevelHistogram(
   int numBits = llvm::Log2_64(numBins);
   int numBitsLaneId = llvm::Log2_64(numThreadPerWarp);
   unsigned numElementsPerThreads = getTotalElemsPerThread(srcType);
-  unsigned numThreadWithUniqueData = getThreadsPerWarp(srcType)[0];
   // The histogram is distributed across threads, each thread owns `numBins /
   // numThreadPerWarp` bins.
   SmallVector<Value> warpLevelHistogram(numBins / numThreadPerWarp, zero);
@@ -43,10 +42,6 @@ static SmallVector<Value> computeWarpLevelHistogram(
         numThreadPerWarp == 32 ? 0xFFFFFFFF : 0xFFFFFFFFFFFFFFFF;
     Value fullMask = b.int_val(numThreadPerWarp, fullMaskValue);
     Value mask = fullMask;
-    // If not all threads have unique data, mask out the redundant ones.
-    if (numThreadWithUniqueData < numThreadPerWarp) {
-      mask = b.int_val(numThreadPerWarp, (1ULL << numThreadWithUniqueData) - 1);
-    }
     for (int i = 0; i < numBitsLaneId; i++) {
       Value updateMask =
           b.select(b.icmp_ne(b.and_(threadId, b.i32_val(1 << i)), zero),
@@ -96,8 +91,6 @@ static SmallVector<Value> computeCrossWarpHistogram(
     Value threadId, int numWarps) {
   auto b = TritonLLVMOpBuilder(loc, rewriter);
   SmallVector<Value> histogramValues;
-  unsigned numWarpsWithUniqueData = mlir::triton::gpu::getWarpsPerCTA(
-      srcType.getEncoding(), srcType.getShape())[0];
   Value laneId = b.and_(threadId, b.i32_val(numThreadPerWarp - 1));
   // Initialize the shared memory with zeros.
   int64_t numElementPerThread =
@@ -112,19 +105,6 @@ static SmallVector<Value> computeCrossWarpHistogram(
   }
   b.barrier();
   Block *afterAtomics = nullptr;
-  // If some warps have replicated data we need to skip those warps when
-  // accumulating.
-  if (numWarpsWithUniqueData < numWarps) {
-    Block *currentBlock = rewriter.getInsertionBlock();
-    afterAtomics =
-        rewriter.splitBlock(currentBlock, rewriter.getInsertionPoint());
-    Block *atomicBlock = rewriter.createBlock(afterAtomics);
-    rewriter.setInsertionPointToEnd(currentBlock);
-    Value cond = b.icmp_ult(
-        threadId, b.i32_val(numWarpsWithUniqueData * numThreadPerWarp));
-    rewriter.create<LLVM::CondBrOp>(loc, cond, atomicBlock, afterAtomics);
-    rewriter.setInsertionPointToStart(atomicBlock);
-  }
   // Apply atomic add to update the histogram in shared memory.
   for (int i = 0; i < warpLevelHistogram.size(); ++i) {
     Value warpLevelHistogramValue = warpLevelHistogram[i];
@@ -209,6 +189,24 @@ struct HistogramOpConversion
         loc, rewriter, srcType, baseSharedMemPtr, warpLevelHistogram, numBins,
         numThreadsPerWarp, innerDimIndices, threadId, numWarps);
 
+    // Depending on the layout, some threads may have duplicate data. We can
+    // account for this by calculating a "replication factor" and dividing the
+    // results by it to avoid overcounting.
+    auto replicationFactor = numWarps * numThreadsPerWarp;
+    auto threadsPerWarp = getThreadsPerWarp(srcType);
+    auto warpsPerCTA =
+        getWarpsPerCTA(srcType.getEncoding(), srcType.getShape());
+    replicationFactor /= std::accumulate(
+        threadsPerWarp.begin(), threadsPerWarp.end(), 1, std::multiplies<>());
+    replicationFactor /= std::accumulate(warpsPerCTA.begin(), warpsPerCTA.end(),
+                                         1, std::multiplies<>());
+
+    auto b = TritonLLVMOpBuilder(loc, rewriter);
+    for (auto i = 0; i < histogramValue.size(); ++i) {
+      histogramValue[i] =
+          b.sdiv(histogramValue[i], b.i32_val(replicationFactor));
+    }
+
     Value results = packLLElements(loc, typeConverter, histogramValue, rewriter,
                                    op.getType());
     rewriter.replaceOp(op, results);

lib/Dialect/Triton/IR/Ops.cpp

@@ -1019,8 +1019,8 @@ OpFoldResult AdvanceOp::fold(FoldAdaptor adaptor) {
 //-- MakeTensorDescOp --
 void MakeTensorDescOp::build(OpBuilder &builder, OperationState &state,
                              Value base, ValueRange shape, ValueRange strides,
-                             ArrayRef<int32_t> blockShape,
-                             bool isSignedInteger) {
+                             ArrayRef<int32_t> blockShape, bool isSignedInteger,
+                             triton::PaddingOption padding) {
   auto ptrTy = dyn_cast<triton::PointerType>(base.getType());
   if (!ptrTy) {
     llvm::report_fatal_error("Expected pointer type");
@@ -1030,7 +1030,8 @@ void MakeTensorDescOp::build(OpBuilder &builder, OperationState &state,
   auto blockTy = RankedTensorType::get(blockShape64, elemTy);
   auto descTy =
       TensorDescType::get(builder.getContext(), blockTy, isSignedInteger);
-  return build(builder, state, descTy, base, shape, strides);
+  auto paddingAttr = PaddingOptionAttr::get(builder.getContext(), padding);
+  return build(builder, state, descTy, base, shape, strides, paddingAttr);
 }
 
 // The following ops, including `call`, `func`, and `return` are copied and

lib/Dialect/Triton/Transforms/RewriteTensorDescriptorToPointer.cpp

@@ -59,18 +59,21 @@ struct Descriptor {
   Value base;
   ValueRange shape;
   ValueRange strides;
+  Value paddingOption;
 };
 
 Descriptor unpackDescriptor(TensorDescType type, ValueRange pack) {
   int rank = type.getBlockType().getRank();
-  assert(pack.size() == 1 + 2 * static_cast<size_t>(rank) &&
+  assert(pack.size() == 1 + 2 * static_cast<size_t>(rank) + 1 &&
          "Expected tensor descriptors to consist of a pointer, "
-         "followed by 'rank' shape values and 'rank' stride values.");
+         "followed by 'rank' shape values and 'rank' stride values, "
+         "followed by a padding option value.");
 
   Descriptor res;
   res.base = pack[0];
   res.shape = pack.slice(1, rank);
   res.strides = pack.slice(1 + rank, rank);
+  res.paddingOption = pack[1 + 2 * rank];
   return res;
 }
 
@@ -211,16 +214,30 @@ Value generateMask(OpBuilder &builder, const Location &loc,
 }
 
 Value generateOther(OpBuilder &builder, Location loc, Type scalarTy,
-                    ArrayRef<int64_t> blockShape) {
+                    ArrayRef<int64_t> blockShape,
+                    Value paddingOption = nullptr) {
   auto blockTy = RankedTensorType::get(blockShape, scalarTy);
-  auto attr = builder.getZeroAttr(blockTy);
-  return builder.create<arith::ConstantOp>(loc, attr);
+  if (paddingOption && mlir::isa<FloatType>(scalarTy)) {
+    auto floatTy = mlir::cast<FloatType>(scalarTy);
+    auto nan = llvm::APFloat::getNaN(floatTy.getFloatSemantics());
+    auto nanValue = builder.create<arith::ConstantOp>(
+        loc,
+        SplatElementsAttr::get(blockTy, builder.getFloatAttr(floatTy, nan)));
+    auto zeroValue = builder.create<arith::ConstantOp>(
+        loc, SplatElementsAttr::get(blockTy, builder.getZeroAttr(floatTy)));
+    return builder.create<mlir::arith::SelectOp>(loc, paddingOption, nanValue,
+                                                 zeroValue);
+  } else {
+    auto attr = builder.getZeroAttr(blockTy);
+    return builder.create<arith::ConstantOp>(loc, attr);
+  }
 }
 
-Value generateOther(OpBuilder &builder, Location loc, TensorDescType descTy) {
+Value generateOther(OpBuilder &builder, Location loc, TensorDescType descTy,
+                    Value paddingOption = nullptr) {
   auto blockTy = descTy.getSignlessBlockType();
   return generateOther(builder, loc, blockTy.getElementType(),
-                       blockTy.getShape());
+                       blockTy.getShape(), paddingOption);
 }
 
 SmallVector<mlir::Value> castToI64(OpBuilder &builder,
@@ -237,12 +254,17 @@ struct RewriteMakeTensorDesc : OpConversionPattern<triton::MakeTensorDescOp> {
   llvm::LogicalResult
   matchAndRewrite(triton::MakeTensorDescOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    SmallVector<mlir::Value> ptrShapeStrides;
-    llvm::append_values(ptrShapeStrides, adaptor.getBase());
-    llvm::append_range(ptrShapeStrides,
+    SmallVector<mlir::Value> ptrShapeStridesPaddingOption;
+    llvm::append_values(ptrShapeStridesPaddingOption, adaptor.getBase());
+    llvm::append_range(ptrShapeStridesPaddingOption,
                        castToI64(rewriter, adaptor.getShape()));
-    llvm::append_range(ptrShapeStrides, adaptor.getStrides());
-    rewriter.replaceOpWithMultiple(op, {ptrShapeStrides});
+    llvm::append_range(ptrShapeStridesPaddingOption, adaptor.getStrides());
+    auto paddingOption = rewriter.create<mlir::arith::ConstantOp>(
+        op.getLoc(), rewriter.getI1Type(),
+        rewriter.getBoolAttr(adaptor.getPadding() ==
+                             triton::PaddingOption::PAD_NAN));
+    llvm::append_values(ptrShapeStridesPaddingOption, paddingOption);
+    rewriter.replaceOpWithMultiple(op, {ptrShapeStridesPaddingOption});
     return mlir::success();
   }
 };
@@ -258,12 +280,11 @@ struct RewriteLoadPattern : OpConversionPattern<triton::DescriptorLoadOp> {
     auto descTy = op.getDesc().getType();
     auto desc = unpackDescriptor(descTy, adaptor.getDesc());
     auto offsets = castToI64(rewriter, op.getIndices());
-
+    auto other = generateOther(rewriter, loc, descTy, desc.paddingOption);
     auto newLoad = rewriter.replaceOpWithNewOp<triton::LoadOp>(
         op, generatePtr(rewriter, loc, blockShape, desc, offsets),
-        generateMask(rewriter, loc, blockShape, desc, offsets),
-        generateOther(rewriter, loc, descTy), triton::CacheModifier::NONE,
-        triton::EvictionPolicy::NORMAL, false);
+        generateMask(rewriter, loc, blockShape, desc, offsets), other,
+        triton::CacheModifier::NONE, triton::EvictionPolicy::NORMAL, false);
     newLoad->setAttrs(filterSegmentSizes(op->getAttrs()));
 
     return llvm::success();
@@ -327,7 +348,7 @@ struct RewriteGatherPattern : OpConversionPattern<triton::DescriptorGatherOp> {
         rewriter, loc, blockShape, desc, op.getXOffsets(), op.getYOffset());
     auto other = generateOther(rewriter, loc,
                                descTy.getSignlessBlockType().getElementType(),
-                               blockShape);
+                               blockShape, desc.paddingOption);
     auto newLoad = rewriter.replaceOpWithNewOp<triton::LoadOp>(
         op, ptr, mask, other, triton::CacheModifier::NONE,
         triton::EvictionPolicy::NORMAL, false);
@@ -471,13 +492,14 @@ class TritonRewriteTensorDescriptorToPointerPass
     converter.addConversion([](mlir::triton::TensorDescType t,
                                llvm::SmallVectorImpl<mlir::Type> &out) {
       // We convert a tensor descriptor into an pointer, and a shape and stride
-      // for each dimension, i.e., we create 1+2*rank values. Note that tensor
-      // descriptors may be signed/unsigned integers whereas pointers should
-      // always be signless.
+      // for each dimension, and padding option. i.e., we create 1+2*rank+1
+      // values. Note that tensor descriptors may be signed/unsigned integers
+      // whereas pointers should always be signless.
       auto tensorType = t.getSignlessBlockType();
       out.push_back(triton::getPointerType(tensorType.getElementType()));
       out.insert(out.end(), 2 * tensorType.getRank(),
                  mlir::IntegerType::get(t.getContext(), 64));
+      out.push_back(mlir::IntegerType::get(t.getContext(), 1));
       return mlir::success();
     });
 

lib/Dialect/TritonGPU/IR/Dialect.cpp

@@ -1339,12 +1339,6 @@ AMDWmmaEncodingAttr::verify(function_ref<mlir::InFlightDiagnostic()> emitError,
   if (version != 1 && version != 2) {
     return emitError() << "WMMA version must be in the [1, 2] range";
   }
-  // Transposed layout is needed for bypassing LDS between multiple dots.
-  // Version 1 tt.dot results and tt.dot operand layouts are different,
-  // therefore we test and support transposed only for version 2.
-  if (version != 2 && isTransposed) {
-    return emitError() << "Transposed WMMA is supported only for version 2";
-  }
   return success();
 }
 
@@ -2125,10 +2119,10 @@ LogicalResult DotOperandEncodingAttr::verify(
   }
 
   if (auto parentAttr = mlir::dyn_cast<AMDWmmaEncodingAttr>(parent)) {
-    if (kWidth != 16 && parentAttr.getVersion() == 1 ||
+    if (kWidth != 8 && kWidth != 16 && parentAttr.getVersion() == 1 ||
         kWidth != 4 && kWidth != 8 && kWidth != 16 &&
             parentAttr.getVersion() == 2)
-      return emitError() << "ttg.dot_op kWidth parameter must be 16 for "
+      return emitError() << "ttg.dot_op kWidth parameter must be 8/16 for "
                             "gfx11 and 4/8/16 for gfx12 (including packed "
                             "cases for `scaled_dot`)";
     return success();