[MLIR] Prevent invalid IR from being passed outside of RemoveDeadValues #121079

parsifal-47 · 2024-12-25T02:39:48Z

This is a follow-up for #119110 and a fix for #118450

RemoveDeadValues used to delete Values and analyzing the IR at the same time, because of that, isMemoryEffectFree got invalid IR with half-deleted linalg.generic operation. This PR separates analysis and cleanup to prevent such situation.

Thank you!

llvmbot · 2024-12-25T02:40:21Z

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Renat Idrisov (parsifal-47)

Changes

This is a follow-up for #119110 and a fix for #118450

RemoveDeadValues used to delete Values and analyzing the IR at the same time, because of that, isMemoryEffectFree got invalid IR with half-deleted linalg.generic operation. This PR separates analysis and cleanup to prevent such situation.

Thank you!

Full diff: https://github.com/llvm/llvm-project/pull/121079.diff

2 Files Affected:

(modified) mlir/lib/Transforms/RemoveDeadValues.cpp (+146-60)
(modified) mlir/test/Transforms/remove-dead-values.mlir (+26)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 3429008b2f241a..5d4ec66d6905a4 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -72,15 +72,54 @@ using namespace mlir::dataflow;
 
 namespace {
 
+// Set of structures below to be filled with operations and arguments to erase.
+// This is done to separate analysis and tree modification phases,
+// otherwise analysis is operating on half-deleted tree which is incorrect.
+
+struct CleanupFunction {
+  FunctionOpInterface funcOp;
+  BitVector nonLiveArgs;
+  BitVector nonLiveRets;
+};
+
+struct CleanupOperands {
+  Operation *op;
+  BitVector nonLiveOperands;
+};
+
+struct CleanupResults {
+  Operation *op;
+  BitVector nonLiveResults;
+};
+
+struct CleanupBlockArgs {
+  Block *b;
+  BitVector nonLiveArgs;
+};
+
+struct CleanupSuccessorOperands {
+  BranchOpInterface branch;
+  unsigned index;
+  BitVector nonLiveOperands;
+};
+
+struct CleanupList {
+  SmallVector<Operation *> operations;
+  SmallVector<Value> values;
+  SmallVector<CleanupFunction> functions;
+  SmallVector<CleanupOperands> operands;
+  SmallVector<CleanupResults> results;
+  SmallVector<CleanupBlockArgs> blocks;
+  SmallVector<CleanupSuccessorOperands> successorOperands;
+};
+
 // Some helper functions...
 
 /// Return true iff at least one value in `values` is live, given the liveness
 /// information in `la`.
-static bool hasLive(ValueRange values, RunLivenessAnalysis &la) {
+static bool hasLive(ValueRange values, const DenseSet<Value> &deletionSet, RunLivenessAnalysis &la) {
   for (Value value : values) {
-    // If there is a null value, it implies that it was dropped during the
-    // execution of this pass, implying that it was non-live.
-    if (!value)
+    if (deletionSet.contains(value))
       continue;
 
     const Liveness *liveness = la.getLiveness(value);
@@ -92,11 +131,11 @@ static bool hasLive(ValueRange values, RunLivenessAnalysis &la) {
 
 /// Return a BitVector of size `values.size()` where its i-th bit is 1 iff the
 /// i-th value in `values` is live, given the liveness information in `la`.
-static BitVector markLives(ValueRange values, RunLivenessAnalysis &la) {
+static BitVector markLives(ValueRange values, const DenseSet<Value> &deletionSet, RunLivenessAnalysis &la) {
   BitVector lives(values.size(), true);
 
   for (auto [index, value] : llvm::enumerate(values)) {
-    if (!value) {
+    if (deletionSet.contains(value)) {
       lives.reset(index);
       continue;
     }
@@ -115,6 +154,18 @@ static BitVector markLives(ValueRange values, RunLivenessAnalysis &la) {
   return lives;
 }
 
+// DeletionSet is used to track the Values that are scheduled for removal
+void updateDeletionSet(DenseSet<Value> &deletionSet, ValueRange range, const BitVector &nonLive) {
+  for (auto [index, result] : llvm::enumerate(range)) {
+    if (!nonLive[index]) continue;
+    deletionSet.insert(result);
+  }
+}
+
+void updateDeletionSet(DenseSet<Value> &deletionSet, Operation *op, const BitVector &nonLive) {
+  updateDeletionSet(deletionSet, op->getResults(), nonLive);
+}
+
 /// Drop the uses of the i-th result of `op` and then erase it iff toErase[i]
 /// is 1.
 static void dropUsesAndEraseResults(Operation *op, BitVector toErase) {
@@ -174,43 +225,43 @@ static SmallVector<OpOperand *> operandsToOpOperands(OperandRange operands) {
 /// It is assumed that `op` is simple. Here, a simple op is one which isn't a
 /// function-like op, a call-like op, a region branch op, a branch op, a region
 /// branch terminator op, or return-like.
-static void cleanSimpleOp(Operation *op, RunLivenessAnalysis &la) {
-  if (!isMemoryEffectFree(op) || hasLive(op->getResults(), la))
+static void cleanSimpleOp(CleanupList &cl, DenseSet<Value> &deletionSet, Operation *op, RunLivenessAnalysis &la) {
+  if (!isMemoryEffectFree(op) || hasLive(op->getResults(), deletionSet, la))
     return;
 
-  op->dropAllUses();
-  op->erase();
+  cl.operations.push_back(op);
+  updateDeletionSet(deletionSet, op, BitVector(op->getNumResults(), true));
 }
 
 /// Clean a function-like op `funcOp`, given the liveness information in `la`
 /// and the IR in `module`. Here, cleaning means:
 ///   (1) Dropping the uses of its unnecessary (non-live) arguments,
-///   (2) Erasing these arguments,
-///   (3) Erasing their corresponding operands from its callers,
+///   (2) Erasing their corresponding operands from its callers,
+///   (3) Erasing these arguments,
 ///   (4) Erasing its unnecessary terminator operands (return values that are
 ///   non-live across all callers),
 ///   (5) Dropping the uses of these return values from its callers, AND
 ///   (6) Erasing these return values
 /// iff it is not public or external.
-static void cleanFuncOp(FunctionOpInterface funcOp, Operation *module,
+static void cleanFuncOp(CleanupList &cl, DenseSet<Value> &deletionSet,
+                        FunctionOpInterface funcOp, Operation *module,
                         RunLivenessAnalysis &la) {
   if (funcOp.isPublic() || funcOp.isExternal())
     return;
 
   // Get the list of unnecessary (non-live) arguments in `nonLiveArgs`.
   SmallVector<Value> arguments(funcOp.getArguments());
-  BitVector nonLiveArgs = markLives(arguments, la);
+  BitVector nonLiveArgs = markLives(arguments, deletionSet, la);
   nonLiveArgs = nonLiveArgs.flip();
 
   // Do (1).
   for (auto [index, arg] : llvm::enumerate(arguments))
-    if (arg && nonLiveArgs[index])
-      arg.dropAllUses();
+    if (arg && nonLiveArgs[index]) {
+      cl.values.push_back(arg);
+      deletionSet.insert(arg);
+    }
 
   // Do (2).
-  funcOp.eraseArguments(nonLiveArgs);
-
-  // Do (3).
   SymbolTable::UseRange uses = *funcOp.getSymbolUses(module);
   for (SymbolTable::SymbolUse use : uses) {
     Operation *callOp = use.getUser();
@@ -222,7 +273,7 @@ static void cleanFuncOp(FunctionOpInterface funcOp, Operation *module,
         operandsToOpOperands(cast<CallOpInterface>(callOp).getArgOperands());
     for (int index : nonLiveArgs.set_bits())
       nonLiveCallOperands.set(callOpOperands[index]->getOperandNumber());
-    callOp->eraseOperands(nonLiveCallOperands);
+    cl.operands.push_back({callOp, nonLiveCallOperands});
   }
 
   // Get the list of unnecessary terminator operands (return values that are
@@ -253,26 +304,27 @@ static void cleanFuncOp(FunctionOpInterface funcOp, Operation *module,
   for (SymbolTable::SymbolUse use : uses) {
     Operation *callOp = use.getUser();
     assert(isa<CallOpInterface>(callOp) && "expected a call-like user");
-    BitVector liveCallRets = markLives(callOp->getResults(), la);
+    BitVector liveCallRets = markLives(callOp->getResults(), deletionSet, la);
     nonLiveRets &= liveCallRets.flip();
   }
 
-  // Do (4).
+  // Do (3).
   // Note that in the absence of control flow ops forcing the control to go from
   // the entry (first) block to the other blocks, the control never reaches any
   // block other than the entry block, because every block has a terminator.
   for (Block &block : funcOp.getBlocks()) {
     Operation *returnOp = block.getTerminator();
     if (returnOp && returnOp->getNumOperands() == numReturns)
-      returnOp->eraseOperands(nonLiveRets);
+      cl.operands.push_back({returnOp, nonLiveRets});
   }
-  funcOp.eraseResults(nonLiveRets);
+  cl.functions.push_back({funcOp, nonLiveArgs, nonLiveRets});
 
   // Do (5) and (6).
   for (SymbolTable::SymbolUse use : uses) {
     Operation *callOp = use.getUser();
     assert(isa<CallOpInterface>(callOp) && "expected a call-like user");
-    dropUsesAndEraseResults(callOp, nonLiveRets);
+    cl.results.push_back({callOp, nonLiveRets});
+    updateDeletionSet(deletionSet, callOp, nonLiveRets);
   }
 }
 
@@ -297,18 +349,19 @@ static void cleanFuncOp(FunctionOpInterface funcOp, Operation *module,
 /// It is important to note that values in this op flow from operands and
 /// terminator operands (successor operands) to arguments and results (successor
 /// inputs).
-static void cleanRegionBranchOp(RegionBranchOpInterface regionBranchOp,
+static void cleanRegionBranchOp(CleanupList &cl, DenseSet<Value> &deletionSet,
+                                RegionBranchOpInterface regionBranchOp,
                                 RunLivenessAnalysis &la) {
   // Mark live results of `regionBranchOp` in `liveResults`.
   auto markLiveResults = [&](BitVector &liveResults) {
-    liveResults = markLives(regionBranchOp->getResults(), la);
+    liveResults = markLives(regionBranchOp->getResults(), deletionSet, la);
   };
 
   // Mark live arguments in the regions of `regionBranchOp` in `liveArgs`.
   auto markLiveArgs = [&](DenseMap<Region *, BitVector> &liveArgs) {
     for (Region &region : regionBranchOp->getRegions()) {
       SmallVector<Value> arguments(region.front().getArguments());
-      BitVector regionLiveArgs = markLives(arguments, la);
+      BitVector regionLiveArgs = markLives(arguments, deletionSet, la);
       liveArgs[&region] = regionLiveArgs;
     }
   };
@@ -497,9 +550,8 @@ static void cleanRegionBranchOp(RegionBranchOpInterface regionBranchOp,
   // It could never be live because of this op but its liveness could have been
   // attributed to something else.
   if (isMemoryEffectFree(regionBranchOp.getOperation()) &&
-      !hasLive(regionBranchOp->getResults(), la)) {
-    regionBranchOp->dropAllUses();
-    regionBranchOp->erase();
+      !hasLive(regionBranchOp->getResults(), deletionSet, la)) {
+    cl.operations.push_back(regionBranchOp.getOperation());
     return;
   }
 
@@ -538,29 +590,27 @@ static void cleanRegionBranchOp(RegionBranchOpInterface regionBranchOp,
                    terminatorOperandsToKeep);
 
   // Do (1).
-  regionBranchOp->eraseOperands(operandsToKeep.flip());
+  cl.operands.push_back({regionBranchOp, operandsToKeep.flip()});
 
   // Do (2.a) and (2.b).
   for (Region &region : regionBranchOp->getRegions()) {
     assert(!region.empty() && "expected a non-empty region in an op "
                               "implementing `RegionBranchOpInterface`");
-    for (auto [index, arg] : llvm::enumerate(region.front().getArguments())) {
-      if (argsToKeep[&region][index])
-        continue;
-      if (arg)
-        arg.dropAllUses();
-    }
-    region.front().eraseArguments(argsToKeep[&region].flip());
+    BitVector argsToRemove = argsToKeep[&region].flip();
+    cl.blocks.push_back({&region.front(), argsToRemove});
+    updateDeletionSet(deletionSet, region.front().getArguments(), argsToRemove);
   }
 
   // Do (2.c).
   for (Region &region : regionBranchOp->getRegions()) {
     Operation *terminator = region.front().getTerminator();
-    terminator->eraseOperands(terminatorOperandsToKeep[terminator].flip());
+    cl.operands.push_back({terminator, terminatorOperandsToKeep[terminator].flip()});
   }
 
   // Do (3) and (4).
-  dropUsesAndEraseResults(regionBranchOp.getOperation(), resultsToKeep.flip());
+  BitVector resultsToRemove = resultsToKeep.flip();
+  updateDeletionSet(deletionSet, regionBranchOp.getOperation(), resultsToRemove);
+  cl.results.push_back({regionBranchOp.getOperation(), resultsToRemove});
 }
 
 // 1. Iterate over each successor block of the given BranchOpInterface
@@ -572,7 +622,8 @@ static void cleanRegionBranchOp(RegionBranchOpInterface regionBranchOp,
 //    c. Mark each operand as live or dead based on the analysis.
 // 3. Remove dead operands from the branch operation and arguments accordingly
 
-static void cleanBranchOp(BranchOpInterface branchOp, RunLivenessAnalysis &la) {
+static void cleanBranchOp(CleanupList &cl, DenseSet<Value> &deletionSet,
+                          BranchOpInterface branchOp, RunLivenessAnalysis &la) {
   unsigned numSuccessors = branchOp->getNumSuccessors();
 
   // Do (1)
@@ -588,22 +639,53 @@ static void cleanBranchOp(BranchOpInterface branchOp, RunLivenessAnalysis &la) {
       operandValues.push_back(successorOperands[operandIdx]);
     }
 
-    BitVector successorLiveOperands = markLives(operandValues, la);
-
     // Do (3)
-    for (int argIdx = successorLiveOperands.size() - 1; argIdx >= 0; --argIdx) {
-      if (!successorLiveOperands[argIdx]) {
-        if (successorBlock->getNumArguments() < successorOperands.size()) {
-          // if block was cleaned through a different code path
-          // we only need to remove operands from the invokation
-          successorOperands.erase(argIdx);
-          continue;
-        }
+    BitVector successorNonLive = markLives(operandValues, deletionSet, la).flip();
+    updateDeletionSet(deletionSet, successorBlock->getArguments(), successorNonLive);
+    cl.blocks.push_back({successorBlock, successorNonLive});
+    cl.successorOperands.push_back({branchOp, succIdx, successorNonLive});
+  }
+}
+
+void cleanup(CleanupList &cl) {
+  for (auto &op: cl.operations) {
+    op->dropAllUses();
+    op->erase();
+  }
+
+  for (auto &v: cl.values) {
+    v.dropAllUses();
+  }
+
+  for (auto &f: cl.functions) {
+    f.funcOp.eraseArguments(f.nonLiveArgs);
+    f.funcOp.eraseResults(f.nonLiveRets);
+  }
+
+  for (auto &o: cl.operands) {
+    o.op->eraseOperands(o.nonLiveOperands);  }
+
+  for (auto &r: cl.results) {
+    dropUsesAndEraseResults(r.op, r.nonLiveResults);
+  }
 
-        successorBlock->getArgument(argIdx).dropAllUses();
-        successorOperands.erase(argIdx);
-        successorBlock->eraseArgument(argIdx);
-      }
+  for (auto &b: cl.blocks) {
+    // blocks that are accessed via multiple codepaths processed once
+    if (b.b->getNumArguments() != b.nonLiveArgs.size()) continue;
+    for (int i = b.nonLiveArgs.size() - 1; i >= 0; --i) {
+      if (!b.nonLiveArgs[i]) continue;
+      b.b->getArgument(i).dropAllUses();
+      b.b->eraseArgument(i);
+    }
+  }
+  for (auto &op: cl.successorOperands) {
+    SuccessorOperands successorOperands =
+            op.branch.getSuccessorOperands(op.index);
+    // blocks that are accessed via multiple codepaths processed once
+    if (successorOperands.size() != op.nonLiveOperands.size()) continue;
+    for (int i = successorOperands.size() - 1; i >= 0; --i) {
+      if (!op.nonLiveOperands[i]) continue;
+      successorOperands.erase(i);
     }
   }
 }
@@ -616,14 +698,16 @@ struct RemoveDeadValues : public impl::RemoveDeadValuesBase<RemoveDeadValues> {
 void RemoveDeadValues::runOnOperation() {
   auto &la = getAnalysis<RunLivenessAnalysis>();
   Operation *module = getOperation();
+  DenseSet<Value> deletionSet;
+  CleanupList cl;
 
   module->walk([&](Operation *op) {
     if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
-      cleanFuncOp(funcOp, module, la);
+      cleanFuncOp(cl, deletionSet, funcOp, module, la);
     } else if (auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op)) {
-      cleanRegionBranchOp(regionBranchOp, la);
+      cleanRegionBranchOp(cl, deletionSet, regionBranchOp, la);
     } else if (auto branchOp = dyn_cast<BranchOpInterface>(op)) {
-      cleanBranchOp(branchOp, la);
+      cleanBranchOp(cl, deletionSet, branchOp, la);
     } else if (op->hasTrait<::mlir::OpTrait::IsTerminator>()) {
       // Nothing to do here because this is a terminator op and it should be
       // honored with respect to its parent
@@ -631,9 +715,11 @@ void RemoveDeadValues::runOnOperation() {
       // Nothing to do because this op is associated with a function op and gets
       // cleaned when the latter is cleaned.
     } else {
-      cleanSimpleOp(op, la);
+      cleanSimpleOp(cl, deletionSet, op, la);
     }
   });
+
+  cleanup(cl);
 }
 
 std::unique_ptr<Pass> mlir::createRemoveDeadValuesPass() {
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 9273ac01e7ccec..fe7bcbc7c490b6 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -73,6 +73,32 @@ func.func @cleanable_loop_iter_args_value(%arg0: index) -> index {
 
 // -----
 
+// Checking that the arguments of linalg.generic are properly handled
+// All code below is removed as a result of the pass
+//
+#map = affine_map<(d0, d1, d2) -> (0, d1, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+module {
+  func.func @main() {
+    %cst_3 = arith.constant dense<54> : tensor<1x25x13xi32>
+    %cst_7 = arith.constant dense<11> : tensor<1x25x13xi32>
+    // CHECK-NOT: arith.constant
+    %0 = tensor.empty() : tensor<1x25x13xi32>
+    // CHECK-NOT: tensor
+    %1 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%cst_3, %cst_7 : tensor<1x25x13xi32>, tensor<1x25x13xi32>) outs(%0 : tensor<1x25x13xi32>) {
+    // CHECK-NOT: linalg.generic
+    ^bb0(%in: i32, %in_15: i32, %out: i32):
+      %29 = arith.xori %in, %in_15 : i32
+      // CHECK-NOT: arith.xori
+      linalg.yield %29 : i32
+      // CHECK-NOT: linalg.yield
+    } -> tensor<1x25x13xi32>
+    return
+  }
+}
+
+// -----
+
 // Note that this cleanup cannot be done by the `canonicalize` pass.
 //
 // CHECK-LABEL: func.func private @clean_func_op_remove_argument_and_return_value() {

parsifal-47 · 2024-12-25T02:40:27Z

@banach-space please take a look when you get a chance, thank you!

github-actions · 2024-12-25T02:41:28Z

⚠️ We detected that you are using a GitHub private e-mail address to contribute to the repo.
Please turn off Keep my email addresses private setting in your account.
See LLVM Discourse for more information.

github-actions · 2024-12-25T02:43:09Z

✅ With the latest revision this PR passed the C/C++ code formatter.

banach-space

Thanks!

I've quickly skimmed through and left a few comments, but I will need more time for a proper review. This is a bit more involved than I expected - thank you so much for working on this! Main high-level comment:

there should be a very clear distinction between deletionSet and cleanupList.

I've left a suggestion inline.

Otherwise, this looks like the right approach to me :)

mlir/lib/Transforms/RemoveDeadValues.cpp

parsifal-47 · 2024-12-30T22:48:42Z

Thanks!

I've quickly skimmed through and left a few comments, but I will need more time for a proper review. This is a bit more involved than I expected - thank you so much for working on this! Main high-level comment:

there should be a very clear distinction between deletionSet and cleanupList.

I've left a suggestion inline.

Otherwise, this looks like the right approach to me :)

sure, I agree with your comments and will follow-up on resolutions

banach-space · 2024-12-31T15:54:53Z

@joker-eph , when you have some spare time, it would be good if you could confirm that this addresses your concerns raised here:

Check linalg.generic arguments to prevent crashing when they are deleted #119110 (comment)

parsifal-47 · 2024-12-31T16:20:15Z

@banach-space your comments should be resolved, feel free to add more, thank you!

banach-space

Thanks for the updates!

Some more comments inline. Clearly I am struggling with naming - hopefully next year will be better 😅

mlir/lib/Transforms/RemoveDeadValues.cpp

parsifal-47 · 2025-01-01T19:26:04Z

Thanks for the updates!

Some more comments inline. Clearly I am struggling with naming - hopefully next year will be better 😅

no problem, all new comments should be resolved now, thank you!

banach-space

This is becoming much clearer, thank you for all the improvements 🙏🏻

mlir/lib/Transforms/RemoveDeadValues.cpp

banach-space · 2025-01-02T15:59:41Z

mlir/lib/Transforms/RemoveDeadValues.cpp

+/// Process a region branch operation `regionBranchOp` using the liveness
+/// information in `la`. The processing involves two scenarios:
+///
+/// Scenario 1: If the operation has no memory effects and none of its results
+/// are live:
+///   1. Enqueue all its uses for deletion.
+///   2. Enqueue the branch itself for deletion.
+///
+/// Scenario 2: Otherwise:
+///   1. Collect its unnecessary operands (operands forwarded to unnecessary
+///   results or arguments).
+///   2. Process each of its regions.
+///   3. Collect the uses of its unnecessary results (results forwarded from
+///   unnecessary operands
+///      or terminator operands).
+///   4. Add these results to the deletion list.
+///
+/// Processing a region includes:
+///   a. Collecting the uses of its unnecessary arguments (arguments forwarded
+///   from unnecessary operands
+///      or terminator operands).
+///   b. Collecting these unnecessary arguments.
+///   c. Collecting its unnecessary terminator operands (terminator operands
+///   forwarded to unnecessary results
+///      or arguments).
+///
+/// Value Flow Note: In this operation, values flow as follows:
+/// - From operands and terminator operands (successor operands)
+/// - To arguments and results (successor inputs).


I like how you split it into Scenario 1 and Scenario 2, thanks!

I suggest preserving the original format for bullet points ((1) instead of 1.), so that every step can be clearly referred to in the implementation (that's how it's done currently).

Also, I'm wondering how to clearly split steps for Scenario 1 and Scenario 2. How about:

(S1 1) -> (S1 2) for steps 1 and 2 for Scenario 1, and

(S2 1) -> (S2 2) -> (S2 3) -> ... for Scenario 2?

On second thought, given how basic Scenario 1 is, perhaps we only need (1) -> (2) -> (3) for Scenario 2? Either way, it would be good for the steps in comments to match the steps in the code.

I restored original style (1) and also there was (1') and (2') for the short scenario, I got these bullets back, thank you!

banach-space · 2025-01-02T16:12:39Z

Renat, I've added some more reviewers that either have touched or reviewed this code recently (also, CC @codemzs - sadly I couldn't add you as a reviewer).

I am mindful that most of my comments are about ... comments 😅 The code looks good to me, though this is not my area of expertise, hence reaching out to other reviewers.

As for comments, I find them super valuable source of information - thank you for taking the time to document stuff 🙏🏻 And for bearing with me!

Co-authored-by: Andrzej Warzyński <[email protected]>

parsifal-47 · 2025-01-02T18:35:34Z

Renat, I've added some more reviewers that either have touched or reviewed this code recently (also, CC @codemzs - sadly I couldn't add you as a reviewer).

I am mindful that most of my comments are about ... comments 😅 The code looks good to me, though this is not my area of expertise, hence reaching out to other reviewers.

As for comments, I find them super valuable source of information - thank you for taking the time to document stuff 🙏🏻 And for bearing with me!

Sounds good, I agree that the code is changing to the better, I have resolved the latest, thank you!

mlir/lib/Transforms/RemoveDeadValues.cpp

banach-space · 2025-01-02T19:08:49Z

mlir/lib/Transforms/RemoveDeadValues.cpp

+///   (2) Marking their corresponding operands in its callers for removal.
+///   (3) Identifying and enqueueing unnecessary terminator operands
+///      (return values that are non-live across all callers) for removal.
+///   (4) Enqueueing the non-live arguments themselves for removal.


arguments + return vals?

I looked at this description more carefully and it is not (1), (2), and (3), it is 2 and 3 inside of 1, which should not be the case because sublist should have a different numbering. I have removed number from "iterate" step because this is the only step on the top of this small algorithm, and I split 3 in two steps. Please take a look when you have a chance. Thank you!

banach-space · 2025-01-02T19:12:48Z

mlir/lib/Transforms/RemoveDeadValues.cpp

+/// 2. For each successor block, gather all operands from all successors
+///    along with their associated liveness analysis data.


Looks like liveness is actually collected under (3) rather than (2)? Either move code around or update the comment.

yes, comment updated

banach-space · 2025-01-02T19:14:17Z

mlir/lib/Transforms/RemoveDeadValues.cpp

+/// 3. Identify and collect the dead operands from the branch operation
+///    as well as their corresponding arguments.


Should this comment include successors?

cl.successorOperands.push_back({branchOp, succIdx, successorNonLive});

this loop already iterating over successors so successor inside of this loop is singular, if I understood the comment correctly

banach-space · 2025-01-02T19:17:32Z

mlir/lib/Transforms/RemoveDeadValues.cpp

+  }
+}
+
+static void cleanUpDeadVals(RDVFinalCleanupList &RDVFinalCleanupList) {


Please add a comment. To me, something along these lines would be most helpful:

Removes dead values collected in RDVFinalCleanupList. To be run once when all dead values have been collected.

Feel free to re-use.

I'd also be tempted to rename &RDVFinalCleanupList as cleanupList or even list (importantly, note camelCase rather than CamelCase). This hook is clear enough that using simpler names should be fine.

done, thank you!

banach-space · 2025-01-02T19:25:43Z

mlir/lib/Transforms/RemoveDeadValues.cpp

+    for (int i = b.nonLiveArgs.size() - 1; i >= 0; --i) {
+      if (!b.nonLiveArgs[i])
+        continue;
+      b.b->getArgument(i).dropAllUses();
+      b.b->eraseArgument(i);
+    }


Is important to iterate from the end? Also, why not:

Suggested change

for (int i = b.nonLiveArgs.size() - 1; i >= 0; --i) {

if (!b.nonLiveArgs[i])

continue;

b.b->getArgument(i).dropAllUses();

b.b->eraseArgument(i);

}

for (auto [idx, nonLiveArg] : llvm::enumerate(b.nonLiveArgs)) {

if (!nonLiveArg)

continue;

b.b->getArgument(idx).dropAllUses();

b.b->eraseArgument(idx);

}

Or use llvm::foreach.

the collection of arguments is changed while iterating, if we iterate from start, it would be invalid because each deletion is invalidating all successor indexes, that is why iteration is reversed

It's worth adding a note.

sounds good, comment added, thank you!

mlir/lib/Transforms/RemoveDeadValues.cpp

banach-space

LGTM, thank you for improving this, that's much appreciated!

I've left a few minor comments - please address them before merging.

Also, mindful of the holiday season, please wait a couple of days before landing this - just in case folks returning from their breaks want to scan this as well.

mlir/lib/Transforms/RemoveDeadValues.cpp

Co-authored-by: Andrzej Warzyński <[email protected]>

parsifal-47 · 2025-01-05T16:56:09Z

LGTM, thank you for improving this, that's much appreciated!

I've left a few minor comments - please address them before merging.

Also, mindful of the holiday season, please wait a couple of days before landing this - just in case folks returning from their breaks want to scan this as well.

Sure, thanks a lot for your help! I was planning to wait for @joker-eph reaction and also I do not have write permissions so I am unable to merge it anyway :)

parsifal-47 · 2025-01-07T22:16:03Z

@joker-eph could you please take a look when you have a chance? Thank you!

CoTinker · 2025-01-08T01:24:42Z

mlir/lib/Transforms/RemoveDeadValues.cpp

+///      (return values that are non-live across all callers) for removal.
+///   (4) Enqueueing the non-live arguments and return values for removal.
+///   (5) Collecting the uses of these return values in its callers for future
+///   removal.


[nit] you can modify the format slightly.

(5) Collecting the uses of these return values in its callers for future removal.

right, sorry, I forgot about this detail, fixed now, thank you!

CoTinker · 2025-01-08T01:25:18Z

mlir/lib/Transforms/RemoveDeadValues.cpp

+///   results or arguments).
+///   (2) Process each of its regions.
+///   (3) Collect the uses of its unnecessary results (results forwarded from
+///   unnecessary operands


fixed, thank you!

CoTinker · 2025-01-08T01:26:03Z

mlir/lib/Transforms/RemoveDeadValues.cpp

+///      or terminator operands).
+///   (b) Collecting these unnecessary arguments.
+///   (c) Collecting its unnecessary terminator operands (terminator operands
+///   forwarded to unnecessary results


CoTinker · 2025-01-08T01:26:14Z

mlir/lib/Transforms/RemoveDeadValues.cpp

+/// Iterate through each successor block of `branchOp`.
+/// (1) For each successor block, gather all operands from all successors.
+/// (2) Fetch their associated liveness analysis data and collect for future
+/// removal.


parsifal-47 · 2025-01-13T22:35:14Z

@joker-eph could you please take a look if?
besides that all comments are resolved, I would be happy to finalize it and merge

parsifal-47 · 2025-01-17T15:11:27Z

@banach-space should we merge it? what would you recommend? Thank you!

banach-space · 2025-01-19T21:40:20Z

@banach-space should we merge it?

Yes.

From LGTM - How a Patch Is Accepted

Only approval from a single reviewer is required

While it would be great to get feedback from Mehdi, he seems to be offline (I've not seen him interacting online this month). From what I can tell, his concerns from #119110 have been addressed.

So yes, let's land this. This is all within LLVM's "code review" policies.

CoTinker

LGTM, thanks for your work.

parsifal-47 · 2025-01-20T02:39:09Z

@banach-space @CoTinker thank you for your review and patience, please merge it then, I do not have write access to do that myself.

banach-space · 2025-01-20T14:44:59Z

@banach-space @CoTinker thank you for your review and patience, please merge it then, I do not have write access to do that myself.

Thank you for seeing this through and for improving MLIR for all of us 🙏🏻 :) This is much appreciated!

EDIT
Due to this commit: 429ae50, I will be listed as a co-author in the final commit (default behaviour for GitHub's "Squash and merge"). However, 99% of this PR is Renat's work - I've only made a few small suggestions inline.

Prevent invalid IR from being passed outside of RemoveDeadValues

2db7609

llvmbot added mlir:core MLIR Core Infrastructure mlir labels Dec 25, 2024

parsifal-47 changed the title ~~Prevent invalid IR from being passed outside of RemoveDeadValues~~ [MLIR] Prevent invalid IR from being passed outside of RemoveDeadValues Dec 25, 2024

Update formatting

38edd5a

banach-space reviewed Dec 30, 2024

View reviewed changes

banach-space requested a review from joker-eph December 31, 2024 15:53

parsifal-47 added 2 commits December 31, 2024 16:14

Addressing review feedback

ff9a4c1

Formatting

97a219c

banach-space reviewed Dec 31, 2024

View reviewed changes

parsifal-47 added 4 commits January 1, 2025 19:18

Addressing review feedback

3af2d06

Formatting

506a3dc

Addressing review feedback

3aac513

Formatting

398e8db

banach-space reviewed Jan 2, 2025

View reviewed changes

banach-space requested review from srishti-pm and CoTinker January 2, 2025 16:05

parsifal-47 and others added 4 commits January 2, 2025 09:10

Update mlir/lib/Transforms/RemoveDeadValues.cpp

30682b7

Co-authored-by: Andrzej Warzyński <[email protected]>

Update mlir/lib/Transforms/RemoveDeadValues.cpp

6860698

Co-authored-by: Andrzej Warzyński <[email protected]>

Update mlir/lib/Transforms/RemoveDeadValues.cpp

1118697

Co-authored-by: Andrzej Warzyński <[email protected]>

Addressing review feedback

e9dec40

Formatting

58ca561

Formatting

e2ca3db

banach-space reviewed Jan 2, 2025

View reviewed changes

parsifal-47 added 4 commits January 3, 2025 21:13

Addressing review feedback

d60829f

Addressing review feedback

ab30c36

Formatting

68b398a

Formatting

3b8b16c

banach-space approved these changes Jan 5, 2025

View reviewed changes

mlir/lib/Transforms/RemoveDeadValues.cpp Outdated Show resolved Hide resolved

mlir/lib/Transforms/RemoveDeadValues.cpp Outdated Show resolved Hide resolved

mlir/lib/Transforms/RemoveDeadValues.cpp Show resolved Hide resolved

parsifal-47 and others added 3 commits January 5, 2025 08:49

Update mlir/lib/Transforms/RemoveDeadValues.cpp

429ae50

Co-authored-by: Andrzej Warzyński <[email protected]>

Update mlir/lib/Transforms/RemoveDeadValues.cpp

cf7b84f

Co-authored-by: Andrzej Warzyński <[email protected]>

Addressing review feedback

47a8fda

Formatting

7216aa1

CoTinker reviewed Jan 8, 2025

View reviewed changes

Addressing review feedback

ec35781

CoTinker approved these changes Jan 20, 2025

View reviewed changes

banach-space merged commit aa3c31a into llvm:main Jan 20, 2025
8 checks passed

CoTinker mentioned this pull request Feb 20, 2025

[mlir] -remove-dead-values crashes #118450

Closed

		/// 2. For each successor block, gather all operands from all successors
		/// along with their associated liveness analysis data.

		/// 3. Identify and collect the dead operands from the branch operation
		/// as well as their corresponding arguments.

[MLIR] Prevent invalid IR from being passed outside of RemoveDeadValues #121079

[MLIR] Prevent invalid IR from being passed outside of RemoveDeadValues #121079

Uh oh!

Conversation

parsifal-47 commented Dec 25, 2024

Uh oh!

llvmbot commented Dec 25, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

parsifal-47 commented Dec 25, 2024

Uh oh!

github-actions bot commented Dec 25, 2024

Uh oh!

github-actions bot commented Dec 25, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

parsifal-47 commented Dec 30, 2024

Uh oh!

banach-space commented Dec 31, 2024

Uh oh!

parsifal-47 commented Dec 31, 2024

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

parsifal-47 commented Jan 1, 2025

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

banach-space commented Jan 2, 2025

Uh oh!

parsifal-47 commented Jan 2, 2025

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

llvmbot commented Dec 25, 2024 •

edited

Loading

github-actions bot commented Dec 25, 2024 •

edited

Loading