Improve Triton's coalescing pass to support propagation of layout for operations yielding multiple values (#4855)

This PR improves Triton's coalescing pass to support propagation of
layout for operations that yield multiple values. The main enhancement
allows the coalescing algorithm to correctly handle operations with
multiple results by tracking individual values rather than assuming
single-result operations.

Key changes:

- Refactored layout propagation methods to accept specific Value
parameters instead of assuming single results
- Added support for operations with multiple return values in the
coalescing pass
- Enhanced loop handling for scf::ForOp and scf::WhileOp to properly map
individual operands to results

Fixes issues #4854, #4817.

---------

Signed-off-by: Tiotto, Ettore <[email protected]>

Author: etiotto

test/TritonIntelGPU/coalesce.mlir

@@ -558,3 +558,32 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
     tt.return
   }
 }
+
+// -----
+
+// COM: Reproducer for issue #4854
+#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [4, 1], order = [1, 0]}>
+#blocked1 = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [1, 4], order = [1, 0]}>
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+  // CHECK-DAG: [[BLOCKED_LAYOUT:#.*]] = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [4, 1], order = [1, 0]}>
+  // CHECK-DAG: [[BLOCKED_LAYOUT1:#.*]] = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [1, 4], order = [1, 0]}>
+  // CHECK: @test_4854
+  tt.func public @test_4854(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}) {
+    %c0_i32 = arith.constant 0 : i32
+    %c16_i32 = arith.constant 16 : i32
+    %c128_i64 = arith.constant 128 : i64
+    %c1_i64 = arith.constant 1 : i64
+    %c32_i32 = arith.constant 32 : i32
+    %0 = tt.make_tensor_ptr %arg0, [%c128_i64, %c128_i64], [%c1_i64, %c1_i64], [%c0_i32, %c0_i32] {order = array<i32: 0, 1>} : <tensor<128x32xf32, #blocked>>
+    %1 = tt.make_tensor_ptr %arg1, [%c128_i64, %c128_i64], [%c1_i64, %c1_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<32x128xf32, #blocked1>>
+    %2:2 = scf.for %arg2 = %c0_i32 to %c32_i32 step %c32_i32 iter_args(%arg3 = %0, %arg4 = %1) -> (!tt.ptr<tensor<128x32xf32, #blocked>>, !tt.ptr<tensor<32x128xf32, #blocked1>>)  : i32 {
+      %5 = tt.advance %arg4, [%c32_i32, %c0_i32] : <tensor<32x128xf32, #blocked1>>
+      scf.yield %arg3, %5 : !tt.ptr<tensor<128x32xf32, #blocked>>, !tt.ptr<tensor<32x128xf32, #blocked1>>
+    }
+    // CHECK: [[ADV:%.*]] = tt.advance {{.*}} : <tensor<128x32xf32, [[BLOCKED_LAYOUT]]>>
+    %3 = tt.advance %2#0, [%c0_i32, %c16_i32] : <tensor<128x32xf32, #blocked>>
+    // CHECK: [[LOAD:%.*]] = tt.load {{.*}} : !tt.ptr<tensor<32x128xf32, [[BLOCKED_LAYOUT1]]>>
+    %4 = tt.load %1 {boundaryCheck = array<i32: 0>, padding = 1 : i32} : !tt.ptr<tensor<32x128xf32, #blocked1>>
+    tt.return
+  }
+}

third_party/intel/lib/TritonIntelGPUTransforms/Coalesce.cpp

@@ -7,11 +7,13 @@
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Value.h"
 #include "mlir/IR/Verifier.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Support/LLVM.h"
 #include "triton/Dialect/Triton/IR/Types.h"
 #include "triton/Dialect/Triton/IR/Utility.h"
 #include "triton/Dialect/TritonGPU/Transforms/Utility.h"
 #include "triton/Tools/StrUtil.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
@@ -154,55 +156,58 @@ struct CoalescePass
     return false;
   }
 
-  // Change the \p layout of the \p op result and propagate the new result type
-  // to its users.
-  void changeAndPropagateLayout(Operation *op, Attribute layout,
+  // Change the \p layout of the \p op's result \p opRes and propagate the new
+  // result type to its users.
+  void changeAndPropagateLayout(Operation *op, Value opRes, Attribute layout,
                                 IRRewriter &rewriter) const {
     assert(op && op->getNumResults() != 0 &&
            "Expecting operation yielding results");
 
     LLVM_DEBUG({
       llvm::dbgs() << "[" DEBUG_TYPE "]: " << "ChangeAndPropagateLayout for: ";
       op->dumpPretty();
+      llvm::dbgs() << "opRes: ";
+      opRes.printAsOperand(llvm::dbgs(), {});
+      llvm::dbgs() << "\n";
     });
 
     rewriter.modifyOpInPlace(op, [&]() {
-      for (Value res : op->getResults()) {
-        if (!tt::isTensorPointerType(res.getType()))
-          continue;
-
-        auto ptrType = cast<tt::PointerType>(res.getType());
-        auto tensorType = cast<RankedTensorType>(ptrType.getPointeeType());
-        res.setType(tt::PointerType::get(getNewType(tensorType, layout),
+      assert(tt::isTensorPointerType(opRes.getType()));
+      auto ptrType = cast<tt::PointerType>(opRes.getType());
+      auto tensorType = cast<RankedTensorType>(ptrType.getPointeeType());
+      opRes.setType(tt::PointerType::get(getNewType(tensorType, layout),
                                          ptrType.getAddressSpace()));
-      }
     });
 
     LLVM_DEBUG({
       llvm::dbgs() << "[" DEBUG_TYPE "]: Coalesced op: ";
       op->dumpPretty();
     });
 
-    propagateLayout(op, layout, rewriter);
+    for (OpResult res : op->getResults())
+      if (res == opRes)
+        propagateLayout(op, res, layout, rewriter);
   }
 
   // Propagate the layout of the \p root operation's result to its users.
-  void propagateLayout(Operation *root, Attribute layout,
+  void propagateLayout(Operation *op, Value opRes, Attribute layout,
                        IRRewriter &rewriter) const {
-    assert(root->getNumResults() != 0 &&
+    assert(op && op->getNumResults() != 0 &&
            "Expecting an operation yielding a result");
-
-    auto mod = root->getParentOfType<ModuleOp>();
+    assert(opRes &&
+           llvm::any_of(op->getResults(),
+                        [&](OpResult res) { return res == opRes; }) &&
+           "Expecting operation to yield 'opRes'");
 
     LLVM_DEBUG({
-      if (!root->getUsers().empty()) {
+      if (!opRes.getUsers().empty()) {
         llvm::dbgs() << "[" DEBUG_TYPE "]: "
-                     << "Propagate layout to operations using: ";
-        root->dumpPretty();
+                     << "Propagate layout to operations using: " << opRes
+                     << "\n";
       }
     });
 
-    for (Operation *user : root->getUsers()) {
+    for (Operation *user : opRes.getUsers()) {
       if (filterUser(user))
         continue;
 
@@ -212,50 +217,71 @@ struct CoalescePass
       });
 
       if (auto forOp = dyn_cast<scf::ForOp>(user)) {
-        propagateLayoutToArgsAndBody(forOp, root, layout, rewriter);
+        propagateLayoutToArgsAndBody(forOp, opRes, layout, rewriter);
         continue;
       }
       if (auto whileOp = dyn_cast<scf::WhileOp>(user)) {
-        propagateLayoutToArgsAndBody(whileOp, root, layout, rewriter);
+        propagateLayoutToArgsAndBody(whileOp, opRes, layout, rewriter);
         continue;
       }
-
       if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
-        if (auto forOp = yieldOp->getParentOfType<scf::ForOp>())
-          propagateLayoutToLoopResults(forOp, layout, rewriter);
-        if (auto whileOp = yieldOp->getParentOfType<scf::WhileOp>())
-          propagateLayoutToLoopResults(whileOp, layout, rewriter);
-        continue;
+        if (auto loopOp = yieldOp->getParentOfType<LoopLikeOpInterface>()) {
+          for (OpOperand &operand : yieldOp->getOpOperands())
+            if (operand.get() == opRes)
+              propagateLayoutToLoopResults(loopOp, operand.getOperandNumber(),
+                                           layout, rewriter);
+          continue;
+        }
       }
 
       LLVM_DEBUG({
         llvm::dbgs() << "[" DEBUG_TYPE "]: After propagating layout:\n";
-        mod->dumpPretty();
+        op->getParentOfType<ModuleOp>()->dumpPretty();
       });
 
-      changeAndPropagateLayout(user, layout, rewriter);
+      for (OpResult res : user->getResults())
+        changeAndPropagateLayout(user, res, layout, rewriter);
     }
   }
 
   // Propagate the layout of the \p arg block argument to its users.
   void propagateLayout(BlockArgument arg, Attribute layout,
                        IRRewriter &rewriter) const {
+    LLVM_DEBUG({
+      if (!arg.getUsers().empty()) {
+        llvm::dbgs() << "[" DEBUG_TYPE "]: "
+                     << "Propagate layout to operations using: ";
+        arg.printAsOperand(llvm::dbgs(), {});
+        llvm::dbgs() << "\n";
+      }
+    });
+
     for (Operation *user : arg.getUsers()) {
       if (filterUser(user))
         continue;
 
       LLVM_DEBUG({
-        llvm::dbgs() << "[" DEBUG_TYPE "]: " << "arg's user: ";
+        llvm::dbgs() << "[" DEBUG_TYPE "]: " << "user: ";
         user->dumpPretty();
       });
 
-      if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
-        if (auto forOp = yieldOp->getParentOfType<scf::ForOp>())
-          propagateLayoutToLoopResults(forOp, layout, rewriter);
-        if (auto whileOp = yieldOp->getParentOfType<scf::WhileOp>())
-          propagateLayoutToLoopResults(whileOp, layout, rewriter);
+      if (auto forOp = dyn_cast<scf::ForOp>(user)) {
+        propagateLayoutToArgsAndBody(forOp, arg, layout, rewriter);
         continue;
       }
+      if (auto whileOp = dyn_cast<scf::WhileOp>(user)) {
+        propagateLayoutToArgsAndBody(whileOp, arg, layout, rewriter);
+        continue;
+      }
+      if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
+        if (auto loopOp = yieldOp->getParentOfType<LoopLikeOpInterface>()) {
+          for (OpOperand &operand : yieldOp->getOpOperands())
+            if (operand.get() == arg)
+              propagateLayoutToLoopResults(loopOp, operand.getOperandNumber(),
+                                           layout, rewriter);
+          continue;
+        }
+      }
       if (auto condOp = dyn_cast<scf::ConditionOp>(user)) {
         if (auto whileOp = condOp->getParentOfType<scf::WhileOp>()) {
           // Propagate layout to "after" region arguments.
@@ -284,7 +310,8 @@ struct CoalescePass
         continue;
       }
 
-      changeAndPropagateLayout(user, layout, rewriter);
+      for (OpResult res : user->getResults())
+        changeAndPropagateLayout(user, res, layout, rewriter);
     }
 
     LLVM_DEBUG({
@@ -300,74 +327,48 @@ struct CoalescePass
   // loop body that use that argument.
   template <typename OpType, typename = std::enable_if_t<llvm::is_one_of<
                                  OpType, scf::ForOp, scf::WhileOp>::value>>
-  void propagateLayoutToArgsAndBody(OpType loopOp, Operation *root,
+  void propagateLayoutToArgsAndBody(OpType loopOp, Value opRes,
                                     Attribute layout,
                                     IRRewriter &rewriter) const {
-    assert(llvm::any_of(root->getUsers(),
-                        [&](Operation *user) { return user == loopOp; }) &&
-           "Expecting the loop to be a user of the root operation");
-
-    for (BlockArgument arg : loopOp.getRegionIterArgs()) {
-      Value loopArg;
-      if constexpr (std::is_same<OpType, scf::ForOp>::value)
-        loopArg = loopOp.getInitArgs()[arg.getArgNumber() - 1];
-      if constexpr (std::is_same<OpType, scf::WhileOp>::value)
-        loopArg = loopOp.getInits()[arg.getArgNumber()];
-
-      for (OpResult res : root->getResults()) {
-        if (res != loopArg || !tt::isTensorPointerType(res.getType()))
-          continue;
-        // Modify the layout of the loop init argument...
-        tt::PointerType ptrType = cast<tt::PointerType>(arg.getType());
-        auto tensorType = cast<RankedTensorType>(ptrType.getPointeeType());
-        arg.setType(tt::PointerType::get(getNewType(tensorType, layout),
-                                         ptrType.getAddressSpace()));
-        LLVM_DEBUG({
-          llvm::dbgs() << "[" DEBUG_TYPE "]: " << "Propagated layout to: ";
-          arg.printAsOperand(llvm::dbgs(), {});
-          llvm::dbgs() << "\n";
-        });
-
-        // ... and then propagate it to the operations in the loop.
-        propagateLayout(arg, layout, rewriter);
-      }
+    for (auto [initArg, arg] :
+         llvm::zip(loopOp.getInitsMutable(), loopOp.getRegionIterArgs())) {
+      if (initArg.get() != opRes)
+        continue;
+
+      // Modify the layout of the loop init argument...
+      auto ptrType = cast<tt::PointerType>(arg.getType());
+      auto tensorType = cast<RankedTensorType>(ptrType.getPointeeType());
+      arg.setType(tt::PointerType::get(getNewType(tensorType, layout),
+                                       ptrType.getAddressSpace()));
+
+      LLVM_DEBUG({
+        llvm::dbgs() << "[" DEBUG_TYPE "]: " << "Propagated layout to: ";
+        arg.printAsOperand(llvm::dbgs(), {});
+        llvm::dbgs() << "\n";
+      });
+
+      // ... and then propagate it to the operations in the loop.
+      propagateLayout(arg, layout, rewriter);
     }
   }
 
-  // Modify the given loop \p loopOpt and propagate its results to their users.
-  template <typename OpType, typename = std::enable_if_t<llvm::is_one_of<
-                                 OpType, scf::ForOp, scf::WhileOp>::value>>
-  void propagateLayoutToLoopResults(OpType loopOp, Attribute layout,
+  // Modify the \p layout to the loop's operand identified by \p resNum, and
+  // propagate the modified loop results to its users.
+  void propagateLayoutToLoopResults(LoopLikeOpInterface loopOp, unsigned resNum,
+                                    Attribute layout,
                                     IRRewriter &rewriter) const {
-    Operation *yieldOp = nullptr;
-    if constexpr (std::is_same<OpType, scf::ForOp>::value)
-      yieldOp = loopOp.getBody()->getTerminator();
-    if constexpr (std::is_same<OpType, scf::WhileOp>::value)
-      yieldOp = loopOp.getYieldOp();
-
+    Value loopRes = loopOp->getResult(resNum);
     rewriter.modifyOpInPlace(loopOp, [&]() {
-      for (auto [yieldOperandType, res] :
-           llvm::zip(yieldOp->getOperandTypes(), loopOp.getResults())) {
-        Type resType = res.getType();
-        if (yieldOperandType == resType)
-          continue;
-
-        assert(tt::isTensorPointerType(resType) &&
-               tt::isTensorPointerType(yieldOperandType) &&
-               "Expecting blocked pointers");
-        assert(cast<RankedTensorType>(
-                   cast<tt::PointerType>(yieldOperandType).getPointeeType())
-                       .getEncoding() == layout &&
-               "Unexpected layout");
-
-        auto ptrType = cast<tt::PointerType>(res.getType());
-        RankedTensorType tensorType = ttgi::getRankedTensorType(resType);
-        res.setType(tt::PointerType::get(getNewType(tensorType, layout),
-                                         ptrType.getAddressSpace()));
-      }
+      assert(tt::isTensorPointerType(loopRes.getType()) &&
+             "Expecting blocked pointers");
+      Type resType = loopRes.getType();
+      auto ptrType = cast<tt::PointerType>(resType);
+      RankedTensorType tensorType = ttgi::getRankedTensorType(resType);
+      loopRes.setType(tt::PointerType::get(getNewType(tensorType, layout),
+                                           ptrType.getAddressSpace()));
     });
 
-    propagateLayout(loopOp, layout, rewriter);
+    propagateLayout(loopOp, loopRes, layout, rewriter);
   }
 
   void coalesceOp(Attribute encoding, Operation *op) {
@@ -404,7 +405,8 @@ struct CoalescePass
         }
 
         IRRewriter rewriter(builder);
-        changeAndPropagateLayout(*defOp, encoding, rewriter);
+        changeAndPropagateLayout(*defOp, defOp->getResult(), encoding,
+                                 rewriter);
         newArgs.push_back(operand);
       }
     }