[mlir][transforms] Process RegionBranchOp with empty region

[CoTinker]

This PR adds process for RegionBranchOp with empty region, such as 'else' region of scf.if. Fixes #123246.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Longsheng Mou (CoTinker)

Changes

This PR adds process for RegionBranchOp with empty region, such as 'else' region of scf.if. Fixes #123246.

---

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

2 Files Affected:

• (modified) mlir/lib/Transforms/RemoveDeadValues.cpp (+12-2)
• (modified) mlir/test/Transforms/remove-dead-values.mlir (+16)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 3e7a0cca31c772..c20c54551cdf80 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -375,6 +375,8 @@ static void processRegionBranchOp(RegionBranchOpInterface regionBranchOp,
   // Mark live arguments in the regions of `regionBranchOp` in `liveArgs`.
   auto markLiveArgs = [&](DenseMap<Region *, BitVector> &liveArgs) {
     for (Region &region : regionBranchOp->getRegions()) {
+      if (region.empty())
+        continue;
       SmallVector<Value> arguments(region.front().getArguments());
       BitVector regionLiveArgs = markLives(arguments, nonLiveSet, la);
       liveArgs[&region] = regionLiveArgs;
@@ -420,6 +422,8 @@ static void processRegionBranchOp(RegionBranchOpInterface regionBranchOp,
   auto markNonForwardedReturnValues =
       [&](DenseMap<Operation *, BitVector> &nonForwardedRets) {
         for (Region &region : regionBranchOp->getRegions()) {
+          if (region.empty())
+            continue;
           Operation *terminator = region.front().getTerminator();
           nonForwardedRets[terminator] =
               BitVector(terminator->getNumOperands(), true);
@@ -499,6 +503,8 @@ static void processRegionBranchOp(RegionBranchOpInterface regionBranchOp,
         // Recompute `resultsToKeep` and `argsToKeep` based on
         // `terminatorOperandsToKeep`.
         for (Region &region : regionBranchOp->getRegions()) {
+          if (region.empty())
+            continue;
           Operation *terminator = region.front().getTerminator();
           for (const RegionSuccessor &successor : getSuccessors(&region)) {
             Region *successorRegion = successor.getSuccessor();
@@ -547,6 +553,8 @@ static void processRegionBranchOp(RegionBranchOpInterface regionBranchOp,
 
           // Update the terminator operands that need to be kept.
           for (Region &region : regionBranchOp->getRegions()) {
+            if (region.empty())
+              continue;
             updateOperandsOrTerminatorOperandsToKeep(
                 terminatorOperandsToKeep[region.back().getTerminator()],
                 resultsToKeep, argsToKeep, &region);
@@ -611,8 +619,8 @@ static void processRegionBranchOp(RegionBranchOpInterface regionBranchOp,
 
   // Do (2.a) and (2.b).
   for (Region &region : regionBranchOp->getRegions()) {
-    assert(!region.empty() && "expected a non-empty region in an op "
-                              "implementing `RegionBranchOpInterface`");
+    if (region.empty())
+      continue;
     BitVector argsToRemove = argsToKeep[&region].flip();
     cl.blocks.push_back({&region.front(), argsToRemove});
     collectNonLiveValues(nonLiveSet, region.front().getArguments(),
@@ -621,6 +629,8 @@ static void processRegionBranchOp(RegionBranchOpInterface regionBranchOp,
 
   // Do (2.c).
   for (Region &region : regionBranchOp->getRegions()) {
+    if (region.empty())
+      continue;
     Operation *terminator = region.front().getTerminator();
     cl.operands.push_back(
         {terminator, terminatorOperandsToKeep[terminator].flip()});
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index fe7bcbc7c490b6..e549926b904562 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -408,6 +408,22 @@ func.func @main(%arg3 : i32, %arg4 : i1) {
 
 // -----
 
+// The scf.if operation represents an if-then-else construct for conditionally
+// executing two regions of code. The 'the' region has exactly 1 block, and
+// the 'else' region may have 0 or 1 block. This case is to ensure 'else' region
+// with 0 block not crash.
+
+// CHECK-LABEL: func.func @clean_region_branch_op_with_empty_region
+func.func @clean_region_branch_op_with_empty_region(%arg0: i1, %arg1: memref<f32>) {
+  %cst = arith.constant 1.000000e+00 : f32
+  scf.if %arg0 {
+    memref.store %cst, %arg1[] : memref<f32>
+  }
+  return
+}
+
+// -----
+
 #map = affine_map<(d0)[s0, s1] -> (d0 * s0 + s1)>
 func.func @kernel(%arg0: memref<18xf32>) {
   %c1 = arith.constant 1 : index

[python3kgae]

Hi @Mogball, would you mind taking a look at this?

[Mogball]

Thanks for the PR! This seems like a pretty major bug, and I'm surprised that this is the only region affected by this bug, given how prevalent scf.if with an empty else region must be.

That being said, I am not convinced this is the right fix. First, if this indeed an issue, there will be other passes that suffer from the same bug. We should clarify the contract of RegionBranchOpInterface has with its operations. Should ops be allowed to have empty regions? For example, IfOp::getRegionSuccessors will ignore the else region if it's empty.

I personally think scf.if should ban empty else region and simply require it to be else { scf.yield } and we can sugar this in the syntax.

[python3kgae]

Thanks for the PR! This seems like a pretty major bug, and I'm surprised that this is the only region affected by this bug, given how prevalent scf.if with an empty else region must be.

That being said, I am not convinced this is the right fix. First, if this indeed an issue, there will be other passes that suffer from the same bug. We should clarify the contract of RegionBranchOpInterface has with its operations. Should ops be allowed to have empty regions? For example, IfOp::getRegionSuccessors will ignore the else region if it's empty.

I personally think scf.if should ban empty else region and simply require it to be else { scf.yield } and we can sugar this in the syntax.

Thanks for the comment.
Sugar the syntax seems straightforward, we just need to update IfOp::parse.

I'm uncertain about handling the builder, specifically:

void IfOp::build(OpBuilder &builder, OperationState &result,
                 TypeRange resultTypes, Value cond)

This function doesn't create an else region. The caller will use inlineRegionBefore to add the region.
Should we create the else region implicitly? If we do, the callers might need to call eraseBlock before inlineRegionBefore in case they have an elseRegion.
On the other hand, if we don't create the else region implicitly, the caller might not add it at all.

Perhaps we should add a check to ensure the else region is not empty in IfOp::verify() and leave the builder as is?
But that will require fix all the places where the elseRegion is not add.

What do you think?

[Mogball]

Thanks for the PR! This seems like a pretty major bug, and I'm surprised that this is the only region affected by this bug, given how prevalent scf.if with an empty else region must be.
That being said, I am not convinced this is the right fix. First, if this indeed an issue, there will be other passes that suffer from the same bug. We should clarify the contract of RegionBranchOpInterface has with its operations. Should ops be allowed to have empty regions? For example, IfOp::getRegionSuccessors will ignore the else region if it's empty.
I personally think scf.if should ban empty else region and simply require it to be else { scf.yield } and we can sugar this in the syntax.

Thanks for the comment. Sugar the syntax seems straightforward, we just need to update IfOp::parse.

I'm uncertain about handling the builder, specifically:

void IfOp::build(OpBuilder &builder, OperationState &result,
                 TypeRange resultTypes, Value cond)

This function doesn't create an else region. The caller will use inlineRegionBefore to add the region. Should we create the else region implicitly? If we do, the callers might need to call eraseBlock before inlineRegionBefore in case they have an elseRegion. On the other hand, if we don't create the else region implicitly, the caller might not add it at all.

Perhaps we should add a check to ensure the else region is not empty in IfOp::verify() and leave the builder as is? But that will require fix all the places where the elseRegion is not add.

What do you think?

That builder creates the then and else regions but leaves them empty. That's fine, since that's the default behaviour of the builders. I think fixing the verifier to require a non-empty region is a good solution, and then you will have to go about and fix places in the code where the caller doesn't create an scf.yield in the else region. Hopefully it doesn't prove to be too much work, but this is the best way forward IMO.

But I would actually like to invite the opinion of others first, cc @matthias-springer @joker-eph

[joker-eph]

Regions can't be dynamically added, they are fixed after creation, which is why the convention has been for ops like scf.if to create both regions immediately. Adding a block and a yield is "not free" and in general we avoided consuming extra resources when not needed, hence the convention about empty regions.

scf.if isn't the only case, for example a function declaration has an empty region as well. This is even codified in the Parser API somehow with the parseOptionalRegion() API which follows the pattern in general of:

  auto *body = result.addRegion();
  OptionalParseResult parseResult =
      parser.parseOptionalRegion(*body, ...);

If we want to change, I would look at it holistically: it goes beyond just scf.if I believe.

[python3kgae]

Regions can't be dynamically added, they are fixed after creation, which is why the convention has been for ops like scf.if to create both regions immediately. Adding a block and a yield is "not free" and in general we avoided consuming extra resources when not needed, hence the convention about empty regions.

scf.if isn't the only case, for example a function declaration has an empty region as well. This is even codified in the Parser API somehow with the parseOptionalRegion() API which follows the pattern in general of:

  auto *body = result.addRegion();
  OptionalParseResult parseResult =
      parser.parseOptionalRegion(*body, ...);

If we want to change, I would look at it holistically: it goes beyond just scf.if I believe.

So, it is better for RegionBranchOpInterface not to assume a non-empty region when using getRegions, since getRegions is not a method of RegionBranchOpInterface and there're many cases empty regions are allowed right now?

[Mogball]

An alternate solution would be to change the code here to look only at the reachable regions by traversing the region successors. If this happens to be the only pass with this issue, I'm okay with the original change if scf.if containing an empty region is intentional.

[python3kgae]

An alternate solution would be to change the code here to look only at the reachable regions by traversing the region successors. If this happens to be the only pass with this issue, I'm okay with the original change if scf.if containing an empty region is intentional.

I did a quick scan for the use of RegionBranchOpInterface.
When using getRegions on a RegionBranchOpInterface, most cases use for (Block &block : region).
Some use the region in for (Region &region : regionBranchOp->getRegions()) as an argument for other functions, which don’t access inside of the region.

The only place I saw that assumes the region should not be empty is RemoveDeadValues.

[CoTinker]

Thanks for your review. I am currently on vacation and will carefully look into your feedback once I return to the office. Thank you very much!

[CoTinker]

Yes, the main difference lies in using (Region &region : regionBranchOp->getRegions()) versus (Block &block : region). Additionally, if we use getSuccessorRegions, we would need to pass another RegionBranchPoint argument, which complicates access to the region. Therefore, I believe simply verifying that the region is not empty should be sufficient. What do you think?

[CoTinker]

Thanks.