[MLIR] Removing dead values for branches #117501

parsifal-47 · 2024-11-24T18:33:50Z

Fixing RemoveDeadValues to properly remove arguments from BranchOpInterface operations.
This is a follow-up for: #117405
cc: @joker-eph @codemzs

llvmbot · 2024-11-24T18:34:22Z

@llvm/pr-subscribers-mlir-core

Author: Renat Idrisov (parsifal-47)

Changes

Fixing RemoveDeadValues to properly remove arguments from BranchOpInterface operations.
This is a follow-up for: #117405
cc: @joker-eph @codemzs

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

2 Files Affected:

(modified) mlir/lib/Transforms/RemoveDeadValues.cpp (+40-15)
(modified) mlir/test/Transforms/remove-dead-values.mlir (+20-3)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 0aa9dcb36681b3..638726e1212772 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -563,6 +563,44 @@ static void cleanRegionBranchOp(RegionBranchOpInterface regionBranchOp,
   dropUsesAndEraseResults(regionBranchOp.getOperation(), resultsToKeep.flip());
 }
 
+// 1. Iterate over each successor block of the given BranchOpInterface
+// operation.
+// 2. For each successor block:
+//    a. Retrieve the operands passed to the successor.
+//    b. Use the provided liveness analysis (`RunLivenessAnalysis`) to determine
+//    which
+//       operands are live in the successor block.
+//    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) {
+  unsigned numSuccessors = branchOp->getNumSuccessors();
+
+  // Do (1)
+  for (unsigned succIdx = 0; succIdx < numSuccessors; ++succIdx) {
+    Block *successorBlock = branchOp->getSuccessor(succIdx);
+
+    // Do (2)
+    SuccessorOperands successorOperands =
+        branchOp.getSuccessorOperands(succIdx);
+    SmallVector<Value> operandValues;
+    for (unsigned operandIdx = 0; operandIdx < successorOperands.size();
+         ++operandIdx) {
+      operandValues.push_back(successorOperands[operandIdx]);
+    }
+
+    BitVector successorLiveOperands = markLives(operandValues, la);
+
+    // Do (3)
+    for (int argIdx = successorLiveOperands.size() - 1; argIdx >= 0; --argIdx) {
+      if (!successorLiveOperands[argIdx]) {
+        successorOperands.erase(argIdx);
+        successorBlock->eraseArgument(argIdx);
+      }
+    }
+  }
+}
+
 struct RemoveDeadValues : public impl::RemoveDeadValuesBase<RemoveDeadValues> {
   void runOnOperation() override;
 };
@@ -572,26 +610,13 @@ void RemoveDeadValues::runOnOperation() {
   auto &la = getAnalysis<RunLivenessAnalysis>();
   Operation *module = getOperation();
 
-  // The removal of non-live values is performed iff there are no branch ops,
-  // and all symbol user ops present in the IR are call-like.
-  WalkResult acceptableIR = module->walk([&](Operation *op) {
-    if (op == module)
-      return WalkResult::advance();
-    if (isa<BranchOpInterface>(op)) {
-      op->emitError() << "cannot optimize an IR with branch ops\n";
-      return WalkResult::interrupt();
-    }
-    return WalkResult::advance();
-  });
-
-  if (acceptableIR.wasInterrupted())
-    return signalPassFailure();
-
   module->walk([&](Operation *op) {
     if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
       cleanFuncOp(funcOp, module, la);
     } else if (auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op)) {
       cleanRegionBranchOp(regionBranchOp, la);
+    } else if (auto branchOp = dyn_cast<BranchOpInterface>(op)) {
+      cleanBranchOp(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
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 826f6159a36b67..fda7ef3fe673e4 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -28,15 +28,32 @@ module @named_module_acceptable {
 
 // -----
 
-// The IR remains untouched because of the presence of a branch op `cf.cond_br`.
+// The IR is optimized regardless of the presence of a branch op `cf.cond_br`.
 //
-func.func @dont_touch_unacceptable_ir_has_cleanable_simple_op_with_branch_op(%arg0: i1) {
+func.func @acceptable_ir_has_cleanable_simple_op_with_branch_op(%arg0: i1) {
   %non_live = arith.constant 0 : i32
-  // expected-error @+1 {{cannot optimize an IR with branch ops}}
+  // CHECK-NOT: non_live
   cf.cond_br %arg0, ^bb1(%non_live : i32), ^bb2(%non_live : i32)
 ^bb1(%non_live_0 : i32):
+  // CHECK-NOT: non_live_0
   cf.br ^bb3
 ^bb2(%non_live_1 : i32):
+  // CHECK-NOT: non_live_1
+  cf.br ^bb3
+^bb3:
+  return
+}
+
+// -----
+
+// Arguments of unconditional branch op `cf.br` are properly removed.
+//
+func.func @acceptable_ir_has_cleanable_simple_op_with_unconditional_branch_op(%arg0: i1) {
+  %non_live = arith.constant 0 : i32
+  // CHECK-NOT: non_live
+  cf.br ^bb1(%non_live : i32)
+^bb1(%non_live_1 : i32):
+  // CHECK-NOT: non_live_1
   cf.br ^bb3
 ^bb3:
   return

llvmbot · 2024-11-24T18:34:22Z

@llvm/pr-subscribers-mlir

Author: Renat Idrisov (parsifal-47)

Changes

Fixing RemoveDeadValues to properly remove arguments from BranchOpInterface operations.
This is a follow-up for: #117405
cc: @joker-eph @codemzs

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

2 Files Affected:

(modified) mlir/lib/Transforms/RemoveDeadValues.cpp (+40-15)
(modified) mlir/test/Transforms/remove-dead-values.mlir (+20-3)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 0aa9dcb36681b3..638726e1212772 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -563,6 +563,44 @@ static void cleanRegionBranchOp(RegionBranchOpInterface regionBranchOp,
   dropUsesAndEraseResults(regionBranchOp.getOperation(), resultsToKeep.flip());
 }
 
+// 1. Iterate over each successor block of the given BranchOpInterface
+// operation.
+// 2. For each successor block:
+//    a. Retrieve the operands passed to the successor.
+//    b. Use the provided liveness analysis (`RunLivenessAnalysis`) to determine
+//    which
+//       operands are live in the successor block.
+//    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) {
+  unsigned numSuccessors = branchOp->getNumSuccessors();
+
+  // Do (1)
+  for (unsigned succIdx = 0; succIdx < numSuccessors; ++succIdx) {
+    Block *successorBlock = branchOp->getSuccessor(succIdx);
+
+    // Do (2)
+    SuccessorOperands successorOperands =
+        branchOp.getSuccessorOperands(succIdx);
+    SmallVector<Value> operandValues;
+    for (unsigned operandIdx = 0; operandIdx < successorOperands.size();
+         ++operandIdx) {
+      operandValues.push_back(successorOperands[operandIdx]);
+    }
+
+    BitVector successorLiveOperands = markLives(operandValues, la);
+
+    // Do (3)
+    for (int argIdx = successorLiveOperands.size() - 1; argIdx >= 0; --argIdx) {
+      if (!successorLiveOperands[argIdx]) {
+        successorOperands.erase(argIdx);
+        successorBlock->eraseArgument(argIdx);
+      }
+    }
+  }
+}
+
 struct RemoveDeadValues : public impl::RemoveDeadValuesBase<RemoveDeadValues> {
   void runOnOperation() override;
 };
@@ -572,26 +610,13 @@ void RemoveDeadValues::runOnOperation() {
   auto &la = getAnalysis<RunLivenessAnalysis>();
   Operation *module = getOperation();
 
-  // The removal of non-live values is performed iff there are no branch ops,
-  // and all symbol user ops present in the IR are call-like.
-  WalkResult acceptableIR = module->walk([&](Operation *op) {
-    if (op == module)
-      return WalkResult::advance();
-    if (isa<BranchOpInterface>(op)) {
-      op->emitError() << "cannot optimize an IR with branch ops\n";
-      return WalkResult::interrupt();
-    }
-    return WalkResult::advance();
-  });
-
-  if (acceptableIR.wasInterrupted())
-    return signalPassFailure();
-
   module->walk([&](Operation *op) {
     if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
       cleanFuncOp(funcOp, module, la);
     } else if (auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op)) {
       cleanRegionBranchOp(regionBranchOp, la);
+    } else if (auto branchOp = dyn_cast<BranchOpInterface>(op)) {
+      cleanBranchOp(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
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 826f6159a36b67..fda7ef3fe673e4 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -28,15 +28,32 @@ module @named_module_acceptable {
 
 // -----
 
-// The IR remains untouched because of the presence of a branch op `cf.cond_br`.
+// The IR is optimized regardless of the presence of a branch op `cf.cond_br`.
 //
-func.func @dont_touch_unacceptable_ir_has_cleanable_simple_op_with_branch_op(%arg0: i1) {
+func.func @acceptable_ir_has_cleanable_simple_op_with_branch_op(%arg0: i1) {
   %non_live = arith.constant 0 : i32
-  // expected-error @+1 {{cannot optimize an IR with branch ops}}
+  // CHECK-NOT: non_live
   cf.cond_br %arg0, ^bb1(%non_live : i32), ^bb2(%non_live : i32)
 ^bb1(%non_live_0 : i32):
+  // CHECK-NOT: non_live_0
   cf.br ^bb3
 ^bb2(%non_live_1 : i32):
+  // CHECK-NOT: non_live_1
+  cf.br ^bb3
+^bb3:
+  return
+}
+
+// -----
+
+// Arguments of unconditional branch op `cf.br` are properly removed.
+//
+func.func @acceptable_ir_has_cleanable_simple_op_with_unconditional_branch_op(%arg0: i1) {
+  %non_live = arith.constant 0 : i32
+  // CHECK-NOT: non_live
+  cf.br ^bb1(%non_live : i32)
+^bb1(%non_live_1 : i32):
+  // CHECK-NOT: non_live_1
   cf.br ^bb3
 ^bb3:
   return

github-actions · 2024-11-24T18:37:14Z

⚠️ 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.

joker-eph

Can you try to write a test with a scf.for loop that takes some iter_args?

parsifal-47 · 2024-11-24T23:02:21Z

Can you try to write a test with a scf.for loop that takes some iter_args?

Sure, updated with the test: 1663984
@joker-eph please take a look, thank you!

codemzs · 2024-11-24T18:49:08Z

mlir/test/Transforms/remove-dead-values.mlir

+  %non_live = arith.constant 0 : i32
+  // CHECK-NOT: non_live
+  cf.br ^bb1(%non_live : i32)
+^bb1(%non_live_1 : i32):


Can we create a cf.br op inside this branch and then pass the %non_live_1 value as one of it's successor block argument? We are recursively recreating this scenario in one or both of it's successor blocks and I just want to see if alias analysis will clean up the values that refer to the same dead value.

It would also be nice to have a case where one of these dead values is used in the conditions of cf.br of a similar example.

added a sub-branch with a non_live argument, which made me discover and fix a bug in the implementation, thank you!
for the second ask, I am not sure I understand, if the value is used as a condition for control-flow, it is alive, it can't be dead.

@codemzs updated, let me know what do you think

In ^bb1(%non_live_1 : i32): can you create something like %non_live_2 = arith.constant 0 : i32 and then pass it as a second argument to cf.br ^bb3(%non_live_1 : i32, %non_live_2 : i32)

done, please take a look, thank you!

joker-eph · 2024-11-25T11:37:05Z

mlir/test/Transforms/remove-dead-values.mlir

+  %c1 = arith.constant 1 : index
+  %c10 = arith.constant 10 : index
+  %non_live = arith.constant 0 : index
+  // CHECK-NOT: non_live


The names are never propagated, I don't quite see how these check could test anything actually?

I believe you need to actually check that there is a single item_args.

thank you for noticing!
Changed the condition to // CHECK: scf.for %[[ARG_0:.*]] = %c0 to %c10 step %c1 iter_args(%[[ARG_1:.*]] = %arg0)

@joker-eph updated, please let me know if I understood your comment correctly

I still see a // CHECK-NOT: non_live line 53 right now?

sorry, I was thinking it is limited to scf, removed all CHECK-NOT: <name> conditions, replaced them with positive checks, thank you!

codemzs · 2024-11-27T00:20:36Z

mlir/lib/Transforms/RemoveDeadValues.cpp

+//    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) {


I think this needs to be a recursive solution, you can have the same situation in a conditional branch, for example with in a branch you could declare variables and then pass them to the nested conditional branch in this case it won't be part of your initial successor block arguments.

My first idea was to make a recursive function, but I checked other pieces of this pass and I saw no recursion. It looks like it is applied till stable point is reached and no more values getting deleted. In iterative fashion. If doing recursively, we need to check for maximal depth and cyclic dependencies. And do that for any kinds of operations. Do you have an advice @joker-eph ?

This function is called in a walk() which itself will recursively visit the IR.

That said: what about adding a test to cover the case that @codemzs is describing?

Thanks @joker-eph completely forgot this was being invoked from walk(). We are good on that front.

I added such test, but after looking closely I realized that walk and walk<WalkOrder::PostOrder> does not help with BranchOp, probably because they do not have parent-child relation. Sorry, I was wrong, it does not reapply iteratively by itself. I am going to add a recursion in some form and let you know.

I am quite surprised walk() does not traverse nested branch ops, I would also try creating this IR programmatically to ensure you are not missing anything.

it does, the issue is the ordering, I need "inner" branches to be traversed first, that is what PostOrder does, but in case of Branches there is not so much hierarchy. I am going to play with it a bit and get back to you

It turns out to be simpler, I added a test with branching loop passing multiple dead values around to both conditional and unconditional branches. walk works perfectly when all of them are cleaned consistently. Thank you!

parsifal-47 · 2024-11-28T21:29:57Z

@codemzs @joker-eph your comments should be resolved now, thank you!

parsifal-47 · 2024-12-03T03:35:53Z

@codemzs @joker-eph please take a look once you have a chance, I think I addressed all of your comments, thank you!

mlir/lib/Transforms/RemoveDeadValues.cpp

CoTinker · 2024-12-04T06:37:04Z

And please update the comments of cleanSimpleOp

llvm-project/mlir/lib/Transforms/RemoveDeadValues.cpp

Lines 174 to 177 in 59bc03c

    
           /// It is assumed that `op` is simple. Here, a simple op is one which isn't a 
        
           /// symbol op, a symbol-user op, a region branch op, a branch op, a region 
        
           /// branch terminator op, or return-like. 
        
           static void cleanSimpleOp(Operation *op, RunLivenessAnalysis &la) {

parsifal-47 · 2024-12-04T15:10:11Z

And please update the comments of cleanSimpleOp

cleanSimpleOp has not changed, it does not invoke anything additional now, or maybe I do not understand your comment

CoTinker · 2024-12-04T15:25:17Z

Sorry, I misunderstand it. But maybe we should change the comment says symbol op, a symbol-user op, beacuse symbol op except function-like and symbol-user op except call-like are received by this function now.

parsifal-47 · 2024-12-04T15:28:18Z

a symbol-user op

no problem, sounds good, updated the comment, please take a look, thank you!

CoTinker · 2024-12-04T15:41:11Z

Sorry, I should have been a little more explicit.
Please change the original comments from symbol op, a symbol-user op to function-like op, a call-like op.

codemzs · 2024-12-04T16:55:03Z

mlir/test/Transforms/remove-dead-values.mlir

+^bb3(%non_live_2 : i32, %non_live_6 : i32):
+  // CHECK: ^[[BB3]]:
+  cf.cond_br %arg0, ^bb1(%non_live_2 : i32), ^bb4(%non_live_2 : i32)
+  // CHECK: cf.cond_br %arg0, ^[[BB1]], ^[[BB4:bb[0-9]+]]


My understanding is all of the arith.constant in your tests are not alive, hence all of these branch blocks should not have any arguments but why do we see ^[[BB4:bb[0-9]+]] versus ^[[BB1]]?

Correct, the difference between ^[[BB4:bb[0-9]+]] and ^[[BB1]] that the first one is matching the branch name which was generated by optimization, and the second one is using the name. If I understand the question correctly.

codemzs · 2024-12-04T16:56:19Z

mlir/test/Transforms/remove-dead-values.mlir

-  cf.br ^bb3
-^bb3:
+  // CHECK-NOT: arith.constant
+  cf.br ^bb1(%non_live : i32)


lets add a check for cf.br i.e what is expected after this line.

updated, thank you!

This reverts commit 7ebf784.

parsifal-47 · 2024-12-04T18:09:51Z

Sorry, I should have been a little more explicit. Please change the original comments from symbol op, a symbol-user op to function-like op, a call-like op.

sorry about that, updated, thank you!

codemzs

Thank you, @parsifal-47 for making this change!

CoTinker · 2024-12-05T03:50:02Z

Thanks for your work, LGTM. if you have time, you can take a look at this issue #118450.
It's crash due to dropAlluses before erase, result a empty operands linalg op.

parsifal-47 · 2024-12-05T05:00:05Z

Thanks for your work, LGTM. if you have time, you can take a look at this issue #118450. It's crash due to dropAlluses before erase, result a empty operands linalg op.

Sure, l will be happy take a look!
Meanwhile, do you have write access here to merge this PR?
Thank you!

Removing dead values for branches

437db73

llvmbot added mlir:core MLIR Core Infrastructure mlir labels Nov 24, 2024

parsifal-47 mentioned this pull request Nov 24, 2024

Allow SymbolUserOpInterface operators to be used in RemoveDeadValues Pass #117405

Merged

parsifal-47 changed the title ~~Removing dead values for branches~~ [MLIR] Removing dead values for branches Nov 24, 2024

joker-eph reviewed Nov 24, 2024

View reviewed changes

Adding a test with scf.for and iter_args

1663984

codemzs reviewed Nov 25, 2024

View reviewed changes

joker-eph reviewed Nov 25, 2024

View reviewed changes

parsifal-47 added 2 commits November 25, 2024 21:02

Addressing review comment

4617d02

Addressing Code Review Feedback

473fc46

codemzs reviewed Nov 27, 2024

View reviewed changes

parsifal-47 added 4 commits November 28, 2024 20:37

Addressing Code Review Feedback

a84e5d0

Removing redundant test

9e3d7b0

Bettern name for LIT test

4114c7e

Add second argument to one of the branches

97a5531

joker-eph approved these changes Dec 4, 2024

View reviewed changes

CoTinker reviewed Dec 4, 2024

View reviewed changes

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

CoTinker reviewed Dec 4, 2024

View reviewed changes

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

Addressing review feedback

5680ba5

Addressing review feedback

7ebf784

codemzs reviewed Dec 4, 2024

View reviewed changes

parsifal-47 added 2 commits December 4, 2024 18:06

Revert "Addressing review feedback"

b3e9c10

This reverts commit 7ebf784.

Addressing review feedback

2b63c0f

Addressing review feedback

6672caa

codemzs approved these changes Dec 5, 2024

View reviewed changes

CoTinker approved these changes Dec 5, 2024

View reviewed changes

CoTinker merged commit 0629e9e into llvm:main Dec 5, 2024
8 checks passed

[MLIR] Removing dead values for branches #117501

[MLIR] Removing dead values for branches #117501

Uh oh!

Conversation

parsifal-47 commented Nov 24, 2024

Uh oh!

llvmbot commented Nov 24, 2024

Uh oh!

llvmbot commented Nov 24, 2024

Uh oh!

github-actions bot commented Nov 24, 2024

Uh oh!

joker-eph left a comment

Choose a reason for hiding this comment

Uh oh!

parsifal-47 commented Nov 24, 2024

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!

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!

codemzs Nov 27, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

parsifal-47 commented Nov 28, 2024

Uh oh!

parsifal-47 commented Dec 3, 2024

Uh oh!

Uh oh!

Uh oh!

CoTinker commented Dec 4, 2024

Uh oh!

parsifal-47 commented Dec 4, 2024

Uh oh!

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

Uh oh!

Uh oh!

parsifal-47 commented Dec 4, 2024

Uh oh!

CoTinker commented Dec 4, 2024

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!

codemzs Nov 27, 2024 •

edited

Loading

CoTinker commented Dec 4, 2024 •

edited

Loading