[Mlir-commits] [mlir] 0c1a773 - [MLIR] Simplify affine.if having yield values and trivial conditions

Uday Bondhugula llvmlistbot at llvm.org
Wed Jul 7 00:32:58 PDT 2021


Author: Srishti Srivastava
Date: 2021-07-07T13:02:10+05:30
New Revision: 0c1a7730f5372e862150f313df4d9c44757352f4

URL: https://github.com/llvm/llvm-project/commit/0c1a7730f5372e862150f313df4d9c44757352f4
DIFF: https://github.com/llvm/llvm-project/commit/0c1a7730f5372e862150f313df4d9c44757352f4.diff

LOG: [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 <srishti.srivastava at polymagelabs.com>

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

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/IR/AffineOps.cpp
    mlir/test/Dialect/Affine/simplify-affine-structures.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 95a057c484fb5..9fa53a5fcdcf1 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -1905,7 +1905,7 @@ struct SimplifyDeadElse : public OpRewritePattern<AffineIfOp> {
   }
 };
 
-/// 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;
@@ -1913,35 +1913,48 @@ struct AlwaysTrueOrFalseIf : public OpRewritePattern<AffineIfOp> {
   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();
+    };
 
-    IntegerSet conditionSet = op.getIntegerSet();
+    auto isTriviallyTrue = [](IntegerSet iSet) {
+      return (iSet.getNumEqualities() == 1 && iSet.getNumInequalities() == 0 &&
+              iSet.getConstraint(0) == 0);
+    };
+
+    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();
   }
 };
@@ -2051,6 +2064,7 @@ IntegerSet AffineIfOp::getIntegerSet() {
       ->getAttrOfType<IntegerSetAttr>(getConditionAttrName())
       .getValue();
 }
+
 void AffineIfOp::setIntegerSet(IntegerSet newSet) {
   (*this)->setAttr(getConditionAttrName(), IntegerSetAttr::get(newSet));
 }

diff  --git a/mlir/test/Dialect/Affine/simplify-affine-structures.mlir b/mlir/test/Dialect/Affine/simplify-affine-structures.mlir
index d86b5a3628ccc..6d1b806441738 100644
--- a/mlir/test/Dialect/Affine/simplify-affine-structures.mlir
+++ b/mlir/test/Dialect/Affine/simplify-affine-structures.mlir
@@ -368,7 +368,7 @@ func @test_always_false_if_elimination() {
 }
 
 
-// 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 {
@@ -385,16 +385,97 @@ func @test_dimensional_if_elimination() {
   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
 }


        


More information about the Mlir-commits mailing list