Use LivenessAnalysis for dead block arg elimination

Author: neildhar

include/triton/Dialect/TritonGPU/Transforms/Utility.h

@@ -189,8 +189,9 @@ LogicalResult getConvertBackwardSlice(
     std::function<Value(OpOperand &, Attribute)> getExistingConversion =
         nullptr);
 
-// Populate pattern to remove dead cycles in ForOp.
-void populateForOpDeadArgumentElimination(RewritePatternSet &patterns);
+/// Run a dataflow analysis over \p top to identify block arguments to loops
+/// that are dead, and replace their usage with the corresponding init value.
+void runDeadBlockArgumentElimination(Operation *top);
 
 // Convert an \param index to a multi-dim coordinate given \param shape and
 // \param order.

lib/Dialect/Gluon/Transforms/Canonicalize.cpp

@@ -29,6 +29,7 @@ struct Canonicalize : public gluon::impl::GluonCanonicalizeBase<Canonicalize> {
 } // namespace
 
 void Canonicalize::runOnOperation() {
+  runDeadBlockArgumentElimination(getOperation());
   MLIRContext *ctx = &getContext();
   RewritePatternSet patterns(&getContext());
 
@@ -48,7 +49,6 @@ void Canonicalize::runOnOperation() {
   for (mlir::RegisteredOperationName op : ctx->getRegisteredOperationsByDialect(
            cf::ControlFlowDialect::getDialectNamespace()))
     op.getCanonicalizationPatterns(patterns, ctx);
-  populateForOpDeadArgumentElimination(patterns);
 
   // Populate select Triton canonicalization patterns. The important patterns to
   // EXCLUDE are those that modify layouts, especially `ConvertLayoutOp`

lib/Dialect/Gluon/Transforms/SimplifyControlFlow.cpp

@@ -25,6 +25,7 @@ struct SimplifyControlFlow
 } // namespace
 
 void SimplifyControlFlow::runOnOperation() {
+  runDeadBlockArgumentElimination(getOperation());
   MLIRContext *ctx = &getContext();
   RewritePatternSet patterns(&getContext());
 
@@ -39,7 +40,6 @@ void SimplifyControlFlow::runOnOperation() {
   for (mlir::RegisteredOperationName op : ctx->getRegisteredOperationsByDialect(
            cf::ControlFlowDialect::getDialectNamespace()))
     op.getCanonicalizationPatterns(patterns, ctx);
-  populateForOpDeadArgumentElimination(patterns);
 
   GreedyRewriteConfig config;
   // This is intended to run before AutoLayouts are resolved, in which case

lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp

@@ -1685,10 +1685,12 @@ class TritonGPURemoveLayoutConversionsPass
       m.dump();
     });
 
-    // 4. Apply clean up patterns to remove remove dead convert and dead code
+    // 4. Run dataflow analysis to identify dead code before cleanup.
+    runDeadBlockArgumentElimination(m);
+
+    // 5. Apply clean up patterns to remove remove dead convert and dead code
     // generated by the previous transformations.
     RewritePatternSet cleanUpPatterns2(context);
-    populateForOpDeadArgumentElimination(cleanUpPatterns2);
     scf::ForOp::getCanonicalizationPatterns(cleanUpPatterns2, context);
     scf::IfOp::getCanonicalizationPatterns(cleanUpPatterns2, context);
     ConvertLayoutOp::getCanonicalizationPatterns(cleanUpPatterns2, context);

lib/Dialect/TritonGPU/Transforms/Utility.cpp

@@ -2,6 +2,7 @@
 
 #include <fstream>
 
+#include "mlir/Analysis/DataFlow/LivenessAnalysis.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/Dominance.h"
@@ -1241,138 +1242,23 @@ Operation *convertDistributedOpEncoding(Attribute encoding, Operation *op) {
   return newOp;
 }
 
-namespace {
-
-/// Detect dead arguments in scf.for op by assuming all the values are dead and
-/// propagate liveness property.
-struct ForOpDeadArgElimination : public OpRewritePattern<scf::ForOp> {
-  using OpRewritePattern<scf::ForOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(scf::ForOp forOp,
-                                PatternRewriter &rewriter) const final {
-    Block &block = *forOp.getBody();
-    auto yieldOp = cast<scf::YieldOp>(block.getTerminator());
-    // Assume that nothing is live at the beginning and mark values as live
-    // based on uses.
-    DenseSet<Value> aliveValues;
-    SmallVector<Value> queue;
-    // Helper to mark values as live and add them to the queue of value to
-    // propagate if it is the first time we detect the value as live.
-    auto markLive = [&](Value val) {
-      if (!forOp->isAncestor(val.getParentRegion()->getParentOp()))
-        return;
-      if (aliveValues.insert(val).second)
-        queue.push_back(val);
-    };
-    // Mark all yield operands as live if the associated forOp result has any
-    // use.
-    for (auto result : llvm::enumerate(forOp.getResults())) {
-      if (!result.value().use_empty())
-        markLive(yieldOp.getOperand(result.index()));
-    }
-    if (aliveValues.size() == forOp.getNumResults())
-      return failure();
-    // Operations with side-effects are always live. Mark all theirs operands as
-    // live.
-    block.walk([&](Operation *op) {
-      if (!isa<scf::YieldOp, scf::ForOp>(op) && !wouldOpBeTriviallyDead(op)) {
-        for (Value operand : op->getOperands())
-          markLive(operand);
-      }
-    });
-    // Propagate live property until reaching a fixed point.
-    while (!queue.empty()) {
-      Value value = queue.pop_back_val();
-      if (auto nestedFor = value.getDefiningOp<scf::ForOp>()) {
-        auto result = mlir::cast<OpResult>(value);
-        OpOperand &forOperand = *nestedFor.getTiedLoopInit(result);
-        markLive(forOperand.get());
-        auto nestedYieldOp =
-            cast<scf::YieldOp>(nestedFor.getBody()->getTerminator());
-        Value nestedYieldOperand =
-            nestedYieldOp.getOperand(result.getResultNumber());
-        markLive(nestedYieldOperand);
-        continue;
-      }
-      if (auto nestedIf = value.getDefiningOp<scf::IfOp>()) {
-        auto result = mlir::cast<OpResult>(value);
-        // mark condition as live.
-        markLive(nestedIf.getCondition());
-        for (scf::YieldOp nestedYieldOp :
-             {nestedIf.thenYield(), nestedIf.elseYield()}) {
-          Value nestedYieldOperand =
-              nestedYieldOp.getOperand(result.getResultNumber());
-          markLive(nestedYieldOperand);
-        }
-        continue;
-      }
-      if (Operation *def = value.getDefiningOp()) {
-        // TODO: support while ops.
-        if (isa<scf::WhileOp>(def))
-          return failure();
-        for (Value operand : def->getOperands())
-          markLive(operand);
-        continue;
-      }
-      // If an argument block is live then the associated yield operand and
-      // forOp operand are live.
-      auto arg = mlir::cast<BlockArgument>(value);
-      if (auto forOwner = dyn_cast<scf::ForOp>(arg.getOwner()->getParentOp())) {
-        if (arg.getArgNumber() < forOwner.getNumInductionVars())
-          continue;
-        unsigned iterIdx = arg.getArgNumber() - forOwner.getNumInductionVars();
-        Value yieldOperand =
-            forOwner.getBody()->getTerminator()->getOperand(iterIdx);
-        markLive(yieldOperand);
-        markLive(forOwner.getInitArgs()[iterIdx]);
+void runDeadBlockArgumentElimination(Operation *top) {
+  // The op we are running on must not have any results, because the liveness
+  // analysis will not consider their users.
+  assert(top->hasTrait<OpTrait::ZeroResults>() && "op cannot have results");
+  dataflow::RunLivenessAnalysis la{top};
+
+  // Remove users of dead block arguments in loops by replacing them with their
+  // init values.
+  top->walk([&](Operation *op) {
+    if (auto loopLike = dyn_cast<LoopLikeOpInterface>(op)) {
+      for (auto [idx, arg] : llvm::enumerate(loopLike.getRegionIterArgs())) {
+        const auto *liveness = la.getLiveness(arg);
+        if (liveness && !liveness->isLive)
+          arg.replaceAllUsesWith(loopLike.getInits()[idx]);
       }
     }
-    SmallVector<unsigned> deadArg;
-    for (auto yieldOperand : llvm::enumerate(yieldOp->getOperands())) {
-      if (aliveValues.contains(yieldOperand.value()))
-        continue;
-      if (yieldOperand.value() == block.getArgument(yieldOperand.index() + 1))
-        continue;
-
-      // The yield operand might live outside the loop, e.g.
-      //   %init = ...
-      //   %x = ...
-      //   %y = for iter_args(%unused = %init) {
-      //     yield %x
-      //   }
-      //
-      // In this case, the loop returns %x if it runs 1 or more times, and
-      // otherwise it returns %init.  We cowardly refuse to remove this operand
-      // from the yield.  (We could, but we'd need to prove that the loop runs 0
-      // or >=1 times.)
-      //
-      // As a special case, if it doesn't matter whether the loop runs 0 or >=1
-      // times (because the loop returns the same value in both cases) then we
-      // can still mark the operand as dead. This occurs in the above example
-      // when %init is the same as %x.
-      if (!forOp->isAncestor(
-              yieldOperand.value().getParentRegion()->getParentOp()) &&
-          yieldOperand.value() != forOp.getInitArgs()[yieldOperand.index()])
-        continue;
-
-      deadArg.push_back(yieldOperand.index());
-    }
-    bool changed = false;
-    // For simplicity we just replace users of the block arg with init value and
-    // leave the operations and argument removal to dead code elimination.
-    for (unsigned deadArgIdx : deadArg) {
-      BlockArgument arg = block.getArgument(deadArgIdx + 1);
-      changed |= !arg.use_empty();
-      rewriter.replaceAllUsesWith(arg, forOp.getTiedLoopInit(arg)->get());
-    }
-    return success(changed);
-  }
-};
-
-} // namespace
-
-void populateForOpDeadArgumentElimination(RewritePatternSet &patterns) {
-  patterns.add<ForOpDeadArgElimination>(patterns.getContext());
+  });
 }
 
 ttg::LocalAllocOp findShmemAlloc(Value operand) {

test/TritonGPU/combine.mlir

@@ -4056,3 +4056,31 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.thr
     tt.return %3 : tensor<32xf32, #ttg.slice<{dim = 0, parent = #ttg.slice<{dim = 0, parent = #blocked1}>}>>
   }
 }
+
+// -----
+
+#blocked = #ttg.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
+module attributes {"ttg.num-warps" = 4 : i32, "ttg.num-ctas" = 1 : i32, "ttg.threads-per-warp" = 32 : i32} {
+  // CHECK-LABEL: @for_arg_used_in_nested_for_bound
+  tt.func public @for_arg_used_in_nested_for_bound(%arg0: i32, %arg1: i32, %arg2: i32) -> i32 {
+    %c0_i32 = arith.constant 0 : i32
+    %c1_i32 = arith.constant 1 : i32
+    %c128_i32 = arith.constant 128 : i32
+    %c256_i32 = arith.constant 256 : i32
+    // CHECK: scf.for
+    // CHECK-SAME: iter_args(%[[OUTER_ARG0:.*]] = %arg0, %[[OUTER_ARG1:.*]] = %c0_i32)
+    %0:2 = scf.for %iv = %c0_i32 to %c128_i32 step %c1_i32 iter_args(%outer_arg0 = %arg0, %outer_arg1 = %c0_i32) -> (i32, i32) : i32 {
+      %1 = arith.remsi %outer_arg1, %arg2 : i32
+      %2 = arith.addi %1, %c1_i32 : i32
+      %3 = arith.muli %2, %c256_i32 : i32
+      // CHECK: scf.for %[[INNER_IV:.*]] = %c0_i32 to %3 step %c1_i32
+      %inner_result = scf.for %inner_iv = %c0_i32 to %3 step %c1_i32 iter_args(%inner_acc = %outer_arg0) -> (i32) : i32 {
+        %sum = arith.addi %inner_acc, %inner_iv : i32
+        scf.yield %sum : i32
+      }
+      %4 = arith.addi %outer_arg0, %c1_i32 : i32
+      scf.yield %inner_result, %4 : i32, i32
+    }
+    tt.return %0#0 : i32
+  }
+}

third_party/nvidia/hopper/lib/Transforms/WarpSpecialization/WSDataPartition.cpp

@@ -1299,8 +1299,8 @@ static bool doDeepCleanup(triton::FuncOp &funcOp,
     });
 
     // delete block arguments
+    runDeadBlockArgumentElimination(funcOp);
     RewritePatternSet cleanUpPatterns(funcOp.getContext());
-    populateForOpDeadArgumentElimination(cleanUpPatterns);
     scf::ForOp::getCanonicalizationPatterns(cleanUpPatterns,
                                             funcOp.getContext());
     scf::IfOp::getCanonicalizationPatterns(cleanUpPatterns,