[MLIR] Simplify affine.if having yield values and trivial conditions

[srishti-pm]

Summary

When an affine.if operation is returning/yielding results and has a
trivially true or false condition, then its 'then' or 'else' block,
respectively, is promoted to the affine.if's parent block and then, the
affine.if operation is replaced by the correct results/yield values.
Relevant test cases are also added.

Signed-off-by: Srishti Srivastava <[e-mail]>

[bondhugula]

Thanks for adding this simplification.

mlir/test/Dialect/Affine/simplify-affine-structures.mlir

403

3 -> Three

407–411

Can we make this set simpler since the goal here is to test affine.if simplification in the presence of yield values?

[srishti-pm]

[MLIR] Simplify affine.if having yield values and trivial conditions

When an affine.if operation is returning/yielding results and has a
trivially true or false condition, then its 'then' or 'else' block,
respectively, is promoted to the affine.if's parent block and then, the
affine.if operation is replaced by the correct results/yield values.
Relevant test cases are also added.

Signed-off-by: Srishti Srivastava <[e-mail]>

Differential Revision: https://reviews.llvm.org/D105418

[bondhugula]

LGTM! Just a few minor comments.

mlir/test/Dialect/Affine/simplify-affine-structures.mlir

391

oldO1 -> zero? More readable below this way.

406

Nit: Else within backticks.

409–410

For readability, you can name these captured variables like c7, c13, etc. Here and below - this is the typical convention.

469–478

All test cases appear to have block local defined constant ops being yielded. Could you just have one where a value defined outside the affine.if is being yielded and you have that block being promoted?

[srishti-pm]

Modified test cases

[pr4tgpt]

LGTM!

[srishti-pm]

Address comments

[srishti-pm]

@mehdi_amini, could you kindly change my Phabricator username to "srishti-pml", which is also my GitHub username attached to this account now. I am unable to do it myself because when I try, a pop-up comes and it says "Only administrators can change usernames".

[xgupta]

In D105418#2861432, @SrishtiSrivastavaPolyMage wrote:

@mehdi_amini, could you kindly change my Phabricator username to "srishti-pml", which is also my GitHub username attached to this account now. I am unable to do it myself because when I try, a pop-up comes and it says "Only administrators can change usernames".

You should simply create a new account with your @polymagelabs.com email address and get a new username srishti-pml otherwise simply change your GitHub username to SrishtiSrivastavPolyMage if you only want consistency.

[srishti-pm]

In D105418#2866810, @xgupta wrote:

In D105418#2861432, @SrishtiSrivastavaPolyMage wrote:

@mehdi_amini, could you kindly change my Phabricator username to "srishti-pml", which is also my GitHub username attached to this account now. I am unable to do it myself because when I try, a pop-up comes and it says "Only administrators can change usernames".

You should simply create a new account with your @polymagelabs.com email address and get a new username srishti-pml otherwise simply change your GitHub username to SrishtiSrivastavPolyMage if you only want consistency.

Hi, actually I want my username in Phabricator to be srishti-pml. So, when I tried changing my Phabricator username, I found out that I don't have the required permissions.

Now, my Phabricator account is associated with my GitHub account, in which my username was SrishtiSrivastavaPolymage before. So, I changed my GitHub username to srishti-pml, hoping that it would change by Phabricator username. But my Phabricator username didn't change. That is why I am requesting @mehdi_amini to kindly change my Phabricator username to srishti-pml.

[xgupta]

Hi, actually I want my username in Phabricator to be srishti-pml.

I said you to create new account.

So, when I tried changing my Phabricator username, I found out that I don't have the required permissions.

Now, my Phabricator account is associated with my GitHub account, in which my username was SrishtiSrivastavaPolymage before. So, I changed my GitHub username to srishti-pml, hoping that it would change by Phabricator username. But my Phabricator username didn't change. That is why I am requesting @mehdi_amini to kindly change my Phabricator username to srishti-pml.

I have no problem. But actually I think _volunteer_ phabricator administrator are not responsible for changing usernames of 20000+ llvm phab usrers or do they?
It is usually difficult to entertain such request :)

[mehdi_amini]

In D105418#2870990, @xgupta wrote:

Hi, actually I want my username in Phabricator to be srishti-pml.

I said you to create new account.

New accounts have the disadvantage of losing the history of the account. Also if you're using GitHub login, you can't just have a new account (I think?).

So, when I tried changing my Phabricator username, I found out that I don't have the required permissions.

Now, my Phabricator account is associated with my GitHub account, in which my username was SrishtiSrivastavaPolymage before. So, I changed my GitHub username to srishti-pml, hoping that it would change by Phabricator username. But my Phabricator username didn't change. That is why I am requesting @mehdi_amini to kindly change my Phabricator username to srishti-pml.

I have no problem. But actually I think _volunteer_ phabricator administrator are not responsible for changing usernames of 20000+ llvm phab usrers or do they?
It is usually difficult to entertain such request :)

It takes ~10s to do, and it is really not common (it happened a handful of times per year so far).

[xgupta]

Awesome, Thanks!

[srishti-pm]

In D105418#2871674, @mehdi_amini wrote:

In D105418#2870990, @xgupta wrote:

Hi, actually I want my username in Phabricator to be srishti-pml.

I said you to create new account.

New accounts have the disadvantage of losing the history of the account. Also if you're using GitHub login, you can't just have a new account (I think?).

So, when I tried changing my Phabricator username, I found out that I don't have the required permissions.

Now, my Phabricator account is associated with my GitHub account, in which my username was SrishtiSrivastavaPolymage before. So, I changed my GitHub username to srishti-pml, hoping that it would change by Phabricator username. But my Phabricator username didn't change. That is why I am requesting @mehdi_amini to kindly change my Phabricator username to srishti-pml.

I have no problem. But actually I think _volunteer_ phabricator administrator are not responsible for changing usernames of 20000+ llvm phab usrers or do they?
It is usually difficult to entertain such request :)

It takes ~10s to do, and it is really not common (it happened a handful of times per year so far).

Thanks for the help, @mehdi_amini

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

LogicalResult matchAndRewrite(AffineIfOp ifOp,
 
     rewriter.startRootUpdate(ifOp);
     rewriter.eraseBlock(ifOp.getElseBlock());
     rewriter.finalizeRootUpdate(ifOp);
     return success();
   }
 };
 
-/// Removes Affine.If cond if the condition is always true or false in certain
+/// Removes affine.if cond if the condition is always true or false in certain
 /// trivial cases. Promotes the then/else block in the parent operation block.
 struct AlwaysTrueOrFalseIf : public OpRewritePattern<AffineIfOp> {
   using OpRewritePattern<AffineIfOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(AffineIfOp op,
                                 PatternRewriter &rewriter) const override {
 
-    // If affine.if is returning results then don't remove it.
-    // TODO: Similar simplication can be done when affine.if return results.
-    if (op.getNumResults() > 0)
-      return failure();
+    auto isTriviallyFalse = [](IntegerSet iSet) {
+      return iSet.isEmptyIntegerSet();
+    };
+
+    auto isTriviallyTrue = [](IntegerSet iSet) {
+      return (iSet.getNumEqualities() == 1 && iSet.getNumInequalities() == 0 &&
+              iSet.getConstraint(0) == 0);
+    };
 
-    IntegerSet conditionSet = op.getIntegerSet();
+    IntegerSet affineIfConditions = op.getIntegerSet();
     Block *blockToMove;
-    if (conditionSet.isEmptyIntegerSet()) {
-      // If the else region is not there, simply remove the Affine.if
-      // operation.
-      if (!op.hasElse()) {
+    if (isTriviallyFalse(affineIfConditions)) {
+      // The absence, or equivalently, the emptiness of the else region need not
+      // be checked when affine.if is returning results because if an affine.if
+      // operation is returning results, it always has a non-empty else region.
+      if (op.getNumResults() == 0 && !op.hasElse()) {
+        // If the else region is absent, or equivalently, empty, remove the
+        // affine.if operation (which is not returning any results).
         rewriter.eraseOp(op);
         return success();
       }
       blockToMove = op.getElseBlock();
-    } else if (conditionSet.getNumEqualities() == 1 &&
-               conditionSet.getNumInequalities() == 0 &&
-               conditionSet.getConstraint(0) == 0) {
-      // Condition to check for trivially true condition (0==0).
+    } else if (isTriviallyTrue(affineIfConditions)) {
       blockToMove = op.getThenBlock();
     } else {
       return failure();
     }
-    // Remove the terminator from the block as it already exists in parent
-    // block.
-    Operation *blockTerminator = blockToMove->getTerminator();
-    rewriter.eraseOp(blockTerminator);
+    Operation *blockToMoveTerminator = blockToMove->getTerminator();
+    // Promote the "blockToMove" block to the parent operation block between the
+    // prologue and epilogue of "op".
     rewriter.mergeBlockBefore(blockToMove, op);
-    rewriter.eraseOp(op);
+    // Replace the "op" operation with the operands of the
+    // "blockToMoveTerminator" operation. Note that "blockToMoveTerminator" is
+    // the affine.yield operation present in the "blockToMove" block. It has no
+    // operands when affine.if is not returning results and therefore, in that
+    // case, replaceOp just erases "op". When affine.if is not returning
+    // results, the affine.yield operation can be omitted. It gets inserted
+    // implicitly.
+    rewriter.replaceOp(op, blockToMoveTerminator->getOperands());
+    // Erase the "blockToMoveTerminator" operation since it is now in the parent
+    // operation block, which already has its own terminator.
+    rewriter.eraseOp(blockToMoveTerminator);
     return success();
   }
 };
 } // end anonymous namespace.
 
 static LogicalResult verify(AffineIfOp op) {
   // Verify that we have a condition attribute.
   auto conditionAttr =
p.printOptionalAttrDict(op->getAttrs(),
                           /*elidedAttrs=*/op.getConditionAttrName());
 }
 
 IntegerSet AffineIfOp::getIntegerSet() {
   return (*this)
       ->getAttrOfType<IntegerSetAttr>(getConditionAttrName())
       .getValue();
 }
+
 void AffineIfOp::setIntegerSet(IntegerSet newSet) {
   (*this)->setAttr(getConditionAttrName(), IntegerSetAttr::get(newSet));
 }
 
 void AffineIfOp::setConditional(IntegerSet set, ValueRange operands) {
   setIntegerSet(set);
   (*this)->setOperands(operands);
 }

mlir/test/Dialect/Affine/simplify-affine-structures.mlir

affine.for %arg1 = 1 to 100 {
         call @external1() : () -> ()
       }
     }
   }
   return
 }
 
 
-// Testing: Affine.If is not trivially true or false, nothing happens.
+// Testing: affine.if is not trivially true or false, nothing happens.
 // CHECK-LABEL: func @test_dimensional_if_elimination() {
 func @test_dimensional_if_elimination() {
   affine.for %arg0 = 1 to 10 {
     affine.for %arg1 = 1 to 100 {
       // CHECK: affine.if
       // CHECK: } else {
       affine.if affine_set<(d0, d1) : (d0-1 == 0)> (%arg0, %arg1) {
         call @external() : () -> ()
       } else {
         call @external() : () -> ()
       }
     }
   }
   return
 }
 
-// Testing: Affine.If don't get removed if it is returning results.
+// Testing: affine.if gets removed.
 // CHECK-LABEL: func @test_num_results_if_elimination
-func @test_num_results_if_elimination() -> f32 {
-  %zero = constant 0.0 : f32
+func @test_num_results_if_elimination() -> index {
+  // CHECK: %[[zero:.*]] = constant 0 : index
+  %zero = constant 0 : index
+  %0 = affine.if affine_set<() : ()> () -> index {
+    affine.yield %zero : index
+  } else {
+    affine.yield %zero : index
+  }
+  // CHECK-NEXT: return %[[zero]] : index
+  return %0 : index
+}
+
+
+// Three more test functions involving affine.if operations which are
+// returning results:
+
+// Testing: affine.if gets removed. `Else` block get promoted.
+// CHECK-LABEL: func @test_trivially_false_returning_two_results
+// CHECK-SAME: (%[[arg0:.*]]: index)
+func @test_trivially_false_returning_two_results(%arg0: index) -> (index, index) {
+  // CHECK: %[[c7:.*]] = constant 7 : index
+  // CHECK: %[[c13:.*]] = constant 13 : index
+  %c7 = constant 7 : index
+  %c13 = constant 13 : index
+  // CHECK: %[[c2:.*]] = constant 2 : index
+  // CHECK: %[[c3:.*]] = constant 3 : index
+  %res:2 = affine.if affine_set<(d0, d1) : (5 >= 0, -2 >= 0)> (%c7, %c13) -> (index, index) {
+    %c0 = constant 0 : index
+    %c1 = constant 1 : index
+    affine.yield %c0, %c1 : index, index
+  } else {
+    %c2 = constant 2 : index
+    %c3 = constant 3 : index
+    affine.yield %c7, %arg0 : index, index
+  }
+  // CHECK-NEXT: return %[[c7]], %[[arg0]] : index, index
+  return %res#0, %res#1 : index, index
+}
+
+// Testing: affine.if gets removed. `Then` block get promoted.
+// CHECK-LABEL: func @test_trivially_true_returning_five_results
+func @test_trivially_true_returning_five_results() -> (index, index, index, index, index) {
+  // CHECK: %[[c12:.*]] = constant 12 : index
+  // CHECK: %[[c13:.*]] = constant 13 : index
+  %c12 = constant 12 : index
+  %c13 = constant 13 : index
+  // CHECK: %[[c0:.*]] = constant 0 : index
+  // CHECK: %[[c1:.*]] = constant 1 : index
+  // CHECK: %[[c2:.*]] = constant 2 : index
+  // CHECK: %[[c3:.*]] = constant 3 : index
+  // CHECK: %[[c4:.*]] = constant 4 : index
+  %res:5 = affine.if affine_set<(d0, d1) : (1 >= 0, 3 >= 0)>(%c12, %c13) -> (index, index, index, index, index) {
+    %c0 = constant 0 : index
+    %c1 = constant 1 : index
+    %c2 = constant 2 : index
+    %c3 = constant 3 : index
+    %c4 = constant 4 : index
+    affine.yield %c0, %c1, %c2, %c3, %c4 : index, index, index, index, index
+  } else {
+    %c5 = constant 5 : index
+    %c6 = constant 6 : index
+    %c7 = constant 7 : index
+    %c8 = constant 8 : index
+    %c9 = constant 9 : index
+    affine.yield %c5, %c6, %c7, %c8, %c9 : index, index, index, index, index
+  }
+  // CHECK-NEXT: return %[[c0]], %[[c1]], %[[c2]], %[[c3]], %[[c4]] : index, index, index, index, index
+  return %res#0, %res#1, %res#2, %res#3, %res#4 : index, index, index, index, index
+}
+
+// Testing: affine.if doesn't get removed.
+// CHECK-LABEL: func @test_not_trivially_true_or_false_returning_three_results
+func @test_not_trivially_true_or_false_returning_three_results() -> (index, index, index) {
+  // CHECK: %[[c8:.*]] = constant 8 : index
+  // CHECK: %[[c13:.*]] = constant 13 : index
+  %c8 = constant 8 : index
+  %c13 = constant 13 : index
   // CHECK: affine.if
-  %0 = affine.if affine_set<() : ()> () -> f32 {
-    affine.yield %zero : f32
-    // CHECK: else {
+  %res:3 = affine.if affine_set<(d0, d1) : (d0 - 1 == 0)>(%c8, %c13) -> (index, index, index) {
+    %c0 = constant 0 : index
+    %c1 = constant 1 : index
+    %c2 = constant 2 : index
+    affine.yield %c0, %c1, %c2 : index, index, index
+  // CHECK: } else {
   } else {
-    affine.yield %zero : f32
+    %c3 = constant 3 : index
+    %c4 = constant 4 : index
+    %c5 = constant 5 : index
+    affine.yield %c3, %c4, %c5 : index, index, index
   }
-  return %0 : f32
+  return %res#0, %res#1, %res#2 : index, index, index
 }