Merge OpenAI Triton commit 2a86177 (#5021)

This PR change the Triton base from 17be6e1 to 2a86177 (Aug 19).
Pass rate: 98.85%

Author: whitneywhtsang

.github/workflows/integration-tests-amd.yml

@@ -125,7 +125,10 @@ jobs:
           pytest --capture=tee-sys -rfs third_party/amd/python/test/test_extract_slice_concat_op.py
           TRITON_ALWAYS_COMPILE=1 pytest --capture=tee-sys -rfs third_party/amd/python/test/test_scalarize_packed_fops.py
           cd python/test/unit
-          pytest --capture=tee-sys -rfs -n 12 language runtime tools \
+          pytest --capture=tee-sys -rfs -n 12 \
+                 --ignore=blackwell \
+                 --ignore=cuda \
+                 --ignore=instrumentation \
                  --ignore=language/test_line_info.py \
                  --ignore=test_debug.py
           # TODO: uncomment

include/triton/Dialect/Triton/IR/Utility.h

@@ -23,7 +23,7 @@ SmallVector<T> convertType(const VecU &in) {
 }
 
 template <typename Int> Int product(llvm::ArrayRef<Int> arr) {
-  return std::accumulate(arr.begin(), arr.end(), 1, std::multiplies{});
+  return std::accumulate(arr.begin(), arr.end(), 1, std::multiplies<Int>());
 }
 template <typename VecT> auto product(const VecT &vec) {
   return product(llvm::ArrayRef(vec));
@@ -132,8 +132,8 @@ template <typename VecT, typename IdxT>
 // Is `vec` [0, 1, ..., n]?  Returns true on empty list.
 template <typename T> bool isIota(ArrayRef<T> vec) {
   static_assert(std::is_integral_v<T>);
-  for (T i = 0; i < vec.size(); ++i) {
-    if (vec[i] != i) {
+  for (size_t i = 0; i < vec.size(); ++i) {
+    if (vec[i] != static_cast<T>(i)) {
       return false;
     }
   }

include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td

@@ -218,7 +218,7 @@ def TTG_MemDescIndexOp : TTG_Op<"memdesc_index", [Pure, MemDescViewTrait]> {
 
   let results = (outs TTG_MemDescType:$result);
 
-  let assemblyFormat = [{$src `,` $index attr-dict `:` qualified(type($src)) `->` qualified(type($result))}];
+  let assemblyFormat = [{$src `[` $index `]` attr-dict `:` qualified(type($src)) `->` qualified(type($result))}];
 
   let hasVerifier = 1;
 }

lib/Analysis/AxisInfo.cpp

@@ -1261,16 +1261,20 @@ unsigned ModuleAxisInfoAnalysis::getAlignment(Value offsetsValue,
     return 1;
   auto linAttr = gpu::toLinearEncoding(tensorTy);
   auto order = linAttr.getOrder();
-  auto maxMultipleBytes = axisInfo->getDivisibility(order[0]);
-  auto maxContig = axisInfo->getContiguity(order[0]);
 
+  auto divisibility = axisInfo->getDivisibility(order[0]);
   auto elemNumBytes = std::max<unsigned>(elementBitWidth / 8, 1);
-  auto maxMultiple = std::max<int64_t>(maxMultipleBytes / elemNumBytes, 1);
+  auto elemTy = tensorTy.getElementType();
+  auto maxMultiple = isa<PointerType>(elemTy)
+                         ? std::max<int64_t>(divisibility / elemNumBytes, 1)
+                         : divisibility;
+
+  auto maxContig = axisInfo->getContiguity(order[0]);
   unsigned alignment = std::min(maxMultiple, maxContig);
-  LDBG("getAlignment order[0] "
-       << order[0] << " maxMultipleBytes = " << maxMultipleBytes
-       << " maxContig = " << maxContig << " elemNumBits = " << elementBitWidth
-       << " maxMultiple = " << maxMultiple << " alignment " << alignment);
+  LDBG("getAlignment order[0] " << order[0] << " maxContig = " << maxContig
+                                << " elemNumBits = " << elementBitWidth
+                                << " maxMultiple = " << maxMultiple
+                                << " alignment " << alignment);
   LLVM_DEBUG({
     std::string axisStr;
     llvm::raw_string_ostream os(axisStr);

lib/Dialect/Triton/IR/OpInterfaces.cpp

@@ -13,7 +13,7 @@ LogicalResult verifyTransposeOpInterface(Operation *op) {
   TransposeOpInterface transposeOp = cast<TransposeOpInterface>(op);
   auto rank = cast<ShapedType>(transposeOp.getSrc().getType()).getRank();
   auto order = transposeOp.getOrder();
-  if (rank != order.size()) {
+  if (static_cast<size_t>(rank) != order.size()) {
     return op->emitError(
         "order must have the same size as the rank of the operand and result");
   }

lib/Dialect/Triton/IR/Ops.cpp

@@ -465,7 +465,6 @@ static LogicalResult verifyRegionsImpl(Op &op) {
                             << " arguments, but given block with "
                             << block.getNumArguments() << " arguments";
   }
-  unsigned i = 0;
   const auto &blockArgTypes = block.getArgumentTypes();
   for (unsigned i = 0; i < numArgs; ++i) {
     const auto &blockArgTy = blockArgTypes[i];
@@ -966,7 +965,7 @@ LogicalResult BroadcastOp::verify() {
   if (srcShape.size() != resultShape.size()) {
     return emitError("rank of source must be same as rank of result");
   }
-  for (int i = 0; i < srcShape.size(); i++) {
+  for (size_t i = 0; i < srcShape.size(); i++) {
     if (srcShape[i] != 1 && srcShape[i] != resultShape[i]) {
       return emitError("Different dimensions at index ")
              << i << " between source and result.  "
@@ -1288,7 +1287,7 @@ LogicalResult GatherOp::verify() {
   if (getAxis() >= srcTy.getRank()) {
     return emitOpError("gather dimension must be less than the input rank");
   }
-  for (int dim = 0; dim < indicesTy.getRank(); ++dim) {
+  for (uint32_t dim = 0; dim < indicesTy.getRank(); ++dim) {
     if (dim == getAxis())
       continue;
     if (indicesTy.getShape()[dim] != srcTy.getShape()[dim]) {

lib/Dialect/Triton/Transforms/ArithTypeConversion.cpp

@@ -13,7 +13,7 @@ struct RewriteArithSelectOp : mlir::OpConversionPattern<mlir::arith::SelectOp> {
 
   mlir::LogicalResult
   matchAndRewrite(mlir::arith::SelectOp op, OneToNOpAdaptor adaptor,
-                  mlir::ConversionPatternRewriter &rewriter) const {
+                  mlir::ConversionPatternRewriter &rewriter) const override {
     // Note we're replacing the select op with an if op because we are
     // converting one value into many values.
     auto newIf = rewriter.create<mlir::scf::IfOp>(

lib/Dialect/Triton/Transforms/Combine.cpp

@@ -123,7 +123,7 @@ class CombineBroadcastMulReducePattern : public RewritePattern {
       : RewritePattern(ReduceOp::getOperationName(), 1, context) {}
 
   LogicalResult matchAndRewrite(Operation *op,
-                                PatternRewriter &rewriter) const {
+                                PatternRewriter &rewriter) const override {
     auto reduceOp = llvm::dyn_cast<ReduceOp>(op);
     if (!reduceOp)
       return failure();

lib/Dialect/Triton/Transforms/LoopAwareCSE.cpp

@@ -54,8 +54,7 @@ bool LoopCSEDriver::areIterArgsEqual(int i, int j) {
     return false;
   if (llvm::is_contained(argStack, std::make_pair(i, j)))
     return true;
-  BlockArgument aArg = loop.getRegionIterArg(i);
-  BlockArgument bArg = loop.getRegionIterArg(j);
+
   // First, assume the arguments are equal. This is how recursion is broken.
   argStack.push_back({i, j});
   bool result =

lib/Dialect/Triton/Transforms/LoopInvariantCodeMotion.cpp

@@ -55,7 +55,6 @@ class LoopInvariantCodeMotionPass
           [&](Operation *op, Region *) {
             // Create the new mask for load op.
             if (auto loadOp = dyn_cast<LoadOp>(op)) {
-              Value mask = loadOp.getMask();
               IRRewriter rewriter(loopLike);
               Location loc = loopLike->getLoc();
               Value cond;