[MLIR] Simplify affine.if ops with trivial conditions

[shashij-polymagelabs]

Summary

The commit simplifies affine.if ops :
The affine if operation gets removed if the condition is universally true or false and then/else block is merged with the parent block.

Signed-off-by: Shashij Gupta [e-mail]

[pr4tgpt]

Some case styling issues.

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

1912

AffineSet -> affineSet.

1932

block_terminator -> blockTerminator

[bondhugula]

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

1908

Nit: affine.If -> affine.if

1933

You can't use erase() in a rewrite pattern; eraseOp instead.

[bondhugula]

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

1933

rewriter.eraseOp(*block_terminator)

[bondhugula]

This comment was removed by bondhugula.

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

1908

Can you mark this as a TODO? A similar simplification can be done in those cases.

1912

AffineSet -> conditionSet
Avoid generic names.

[shashij-polymagelabs]

Resolved comments

[bondhugula]

Looks good - thanks!

[pr4tgpt]

LGTM!

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

if (ifOp.elseRegion().empty() ||
       return failure();
 
     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
+/// 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();
+
+    IntegerSet conditionSet = op.getIntegerSet();
+    Block *blockToMove;
+    if (conditionSet.isEmptyIntegerSet()) {
+      // If the else region is not there, simply remove the Affine.if
+      // operation.
+      if (!op.hasElse()) {
+        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).
+      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);
+    rewriter.mergeBlockBefore(blockToMove, op);
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
 } // end anonymous namespace.
 
 static LogicalResult verify(AffineIfOp op) {
   // Verify that we have a condition attribute.
   auto conditionAttr =
       op->getAttrOfType<IntegerSetAttr>(op.getConditionAttrName());
   if (!conditionAttr)
     return op.emitOpError(
if (operands.size() < getIntegerSet().getNumInputs() ||
     return success();
   }
 
   return failure();
 }
 
 void AffineIfOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
-  results.add<SimplifyDeadElse>(context);
+  results.add<SimplifyDeadElse, AlwaysTrueOrFalseIf>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // AffineLoadOp
 //===----------------------------------------------------------------------===//
 
 void AffineLoadOp::build(OpBuilder &builder, OperationState &result,
                          AffineMap map, ValueRange operands) {

mlir/test/Dialect/Affine/loop-unswitch.mlir

func @multiple_if(%N : index) {
   }
   affine.for %i = 0 to 100 {
     affine.if affine_set<()[s0] : (s0 >= 0)>()[%N] {
       call @external() : () -> ()
     }
   }
   return
 }
-// CHECK:      affine.if
-// CHECK-NEXT:   call
-// CHECK-NEXT: }
+// CHECK:      call
 // CHECK-NEXT: affine.if
 // CHECK-NEXT:   affine.for
 // CHECK-NEXT:     call
 // CHECK-NEXT:   }
 // CHECK-NEXT: }
 // CHECK-NEXT: return
 
 func private @external()

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

 // RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -simplify-affine-structures | FileCheck %s
 
-// CHECK-DAG: #[[$SET_EMPTY:.*]] = affine_set<() : (1 == 0)>
 // CHECK-DAG: #[[$SET_2D:.*]] = affine_set<(d0, d1) : (d0 - 100 == 0, d1 - 10 == 0, -d0 + 100 >= 0, d1 >= 0)>
 // CHECK-DAG: #[[$SET_7_11:.*]] = affine_set<(d0, d1) : (d0 * 7 + d1 * 5 + 88 == 0, d0 * 5 - d1 * 11 + 60 == 0, d0 * 11 + d1 * 7 - 24 == 0, d0 * 7 + d1 * 5 + 88 == 0)>
 
 // An external function that we will use in bodies to avoid DCE.
 func private @external() -> ()
 
 // CHECK-LABEL: func @test_gaussian_elimination_empty_set0() {
 func @test_gaussian_elimination_empty_set0() {
   affine.for %arg0 = 1 to 10 {
     affine.for %arg1 = 1 to 100 {
-      // CHECK: affine.if #[[$SET_EMPTY]]()
+      // CHECK-NOT: affine.if
       affine.if affine_set<(d0, d1) : (2 == 0)>(%arg0, %arg1) {
         call @external() : () -> ()
       }
     }
   }
   return
 }
 
 // CHECK-LABEL: func @test_gaussian_elimination_empty_set1() {
 func @test_gaussian_elimination_empty_set1() {
   affine.for %arg0 = 1 to 10 {
     affine.for %arg1 = 1 to 100 {
-      // CHECK: affine.if #[[$SET_EMPTY]]()
+      // CHECK-NOT: affine.if
       affine.if affine_set<(d0, d1) : (1 >= 0, -1 >= 0)> (%arg0, %arg1) {
         call @external() : () -> ()
       }
     }
   }
   return
 }
 

 }
 
 // CHECK-LABEL: func @test_gaussian_elimination_empty_set3() {
 func @test_gaussian_elimination_empty_set3() {
   %c7 = constant 7 : index
   %c11 = constant 11 : index
   affine.for %arg0 = 1 to 10 {
     affine.for %arg1 = 1 to 100 {
-      // CHECK: #[[$SET_EMPTY]]()
+      // CHECK-NOT: affine.if
       affine.if affine_set<(d0, d1)[s0, s1] : (d0 - s0 == 0, d0 + s0 == 0, s0 - 1 == 0)>(%arg0, %arg1)[%c7, %c11] {
         call @external() : () -> ()
       }
     }
   }
   return
 }
 
#set_2d_empty = affine_set<(d0, d1)[s0, s1] : (d0 * 7 + d1 * 5 + s0 * 11 + s1 == 0,
                                        d0 - 1 == 0, d0 + 2 == 0)>
 
 // CHECK-LABEL: func @test_gaussian_elimination_empty_set5() {
 func @test_gaussian_elimination_empty_set5() {
   %c7 = constant 7 : index
   %c11 = constant 11 : index
   affine.for %arg0 = 1 to 10 {
     affine.for %arg1 = 1 to 100 {
-      // CHECK: #[[$SET_EMPTY]]()
+      // CHECK-NOT: affine.if
       affine.if #set_2d_empty(%arg0, %arg1)[%c7, %c11] {
         call @external() : () -> ()
       }
     }
   }
   return
 }
 
func @test_fuzz_explosion(%arg0 : index, %arg1 : index, %arg2 : index, %arg3 : index) {
   }
   return
 }
 
 // CHECK-LABEL: func @test_empty_set(%arg0: index) {
 func @test_empty_set(%N : index) {
   affine.for %i = 0 to 10 {
     affine.for %j = 0 to 10 {
-      // CHECK: affine.if #[[$SET_EMPTY]]()
+      // CHECK-NOT: affine.if
       affine.if affine_set<(d0, d1) : (d0 - d1 >= 0, d1 - d0 - 1 >= 0)>(%i, %j) {
         "foo"() : () -> ()
       }
-      // CHECK: affine.if #[[$SET_EMPTY]]()
+      // CHECK-NOT: affine.if
       affine.if affine_set<(d0) : (d0 >= 0, -d0 - 1 >= 0)>(%i) {
         "bar"() : () -> ()
       }
-      // CHECK: affine.if #[[$SET_EMPTY]]()
+      // CHECK-NOT: affine.if
       affine.if affine_set<(d0) : (d0 >= 0, -d0 - 1 >= 0)>(%i) {
         "foo"() : () -> ()
       }
-      // CHECK: affine.if #[[$SET_EMPTY]]()
+      // CHECK-NOT: affine.if
       affine.if affine_set<(d0)[s0, s1] : (d0 >= 0, -d0 + s0 - 1 >= 0, -s0 >= 0)>(%i)[%N, %N] {
         "bar"() : () -> ()
       }
-      // CHECK: affine.if #[[$SET_EMPTY]]()
+      // CHECK-NOT: affine.if
       // The set below implies d0 = d1; so d1 >= d0, but d0 >= d1 + 1.
       affine.if affine_set<(d0, d1, d2) : (d0 - d1 == 0, d2 - d0 >= 0, d0 - d1 - 1 >= 0)>(%i, %j, %N) {
         "foo"() : () -> ()
       }
-      // CHECK: affine.if #[[$SET_EMPTY]]()
+      // CHECK-NOT: affine.if
       // The set below has rational solutions but no integer solutions; GCD test catches it.
       affine.if affine_set<(d0, d1) : (d0*2 -d1*2 - 1 == 0, d0 >= 0, -d0 + 100 >= 0, d1 >= 0, -d1 + 100 >= 0)>(%i, %j) {
         "foo"() : () -> ()
       }
-      // CHECK: affine.if #[[$SET_EMPTY]]()
+      // CHECK-NOT: affine.if
       affine.if affine_set<(d0, d1) : (d1 == 0, d0 - 1 >= 0, - d0 - 1 >= 0)>(%i, %j) {
         "foo"() : () -> ()
       }
     }
   }
   // The tests below test GCDTightenInequalities().
   affine.for %k = 0 to 10 {
     affine.for %l = 0 to 10 {
       // Empty because no multiple of 8 lies between 4 and 7.
-      // CHECK: affine.if #[[$SET_EMPTY]]()
+      // CHECK-NOT: affine.if
       affine.if affine_set<(d0) : (8*d0 - 4 >= 0, -8*d0 + 7 >= 0)>(%k) {
         "foo"() : () -> ()
       }
       // Same as above but with equalities and inequalities.
-      // CHECK: affine.if #[[$SET_EMPTY]]()
+      // CHECK-NOT: affine.if
       affine.if affine_set<(d0, d1) : (d0 - 4*d1 == 0, 4*d1 - 5 >= 0, -4*d1 + 7 >= 0)>(%k, %l) {
         "foo"() : () -> ()
       }
       // Same as above but with a combination of multiple identifiers. 4*d0 +
       // 8*d1 here is a multiple of 4, and so can't lie between 9 and 11. GCD
       // tightening will tighten constraints to 4*d0 + 8*d1 >= 12 and 4*d0 +
       // 8*d1 <= 8; hence infeasible.
-      // CHECK: affine.if #[[$SET_EMPTY]]()
+      // CHECK-NOT: affine.if
       affine.if affine_set<(d0, d1) : (4*d0 + 8*d1 - 9 >= 0, -4*d0 - 8*d1 + 11 >= 0)>(%k, %l) {
         "foo"() : () -> ()
       }
       // Same as above but with equalities added into the mix.
-      // CHECK: affine.if #[[$SET_EMPTY]]()
+      // CHECK-NOT: affine.if
       affine.if affine_set<(d0, d1, d2) : (d0 - 4*d2 == 0, d0 + 8*d1 - 9 >= 0, -d0 - 8*d1 + 11 >= 0)>(%k, %k, %l) {
         "foo"() : () -> ()
       }
     }
   }
 
   affine.for %m = 0 to 10 {
-    // CHECK: affine.if #[[$SET_EMPTY]]()
+    // CHECK-NOT: affine.if
     affine.if affine_set<(d0) : (d0 mod 2 - 3 == 0)> (%m) {
       "foo"() : () -> ()
     }
   }
 
   return
 }
 
 // -----
 
 // An external function that we will use in bodies to avoid DCE.
 func private @external() -> ()
 
 // CHECK-DAG: #[[$SET:.*]] = affine_set<()[s0] : (s0 >= 0, -s0 + 50 >= 0)
-// CHECK-DAG: #[[$EMPTY_SET:.*]] = affine_set<() : (1 == 0)
-// CHECK-DAG: #[[$UNIV_SET:.*]] = affine_set<() : (0 == 0)
 
 // CHECK-LABEL: func @simplify_set
 func @simplify_set(%a : index, %b : index) {
   // CHECK: affine.if #[[$SET]]
   affine.if affine_set<(d0, d1) : (d0 - d1 + d1 + d0 >= 0, 2 >= 0, d0 >= 0, -d0 + 50 >= 0, -d0 + 100 >= 0)>(%a, %b) {
     call @external() : () -> ()
   }
-  // CHECK: affine.if #[[$EMPTY_SET]]
+  // CHECK-NOT: affine.if
   affine.if affine_set<(d0, d1) : (d0 mod 2 - 1 == 0, d0 - 2 * (d0 floordiv 2) == 0)>(%a, %b) {
     call @external() : () -> ()
   }
-  // CHECK: affine.if #[[$UNIV_SET]]
+  // CHECK-NOT: affine.if
   affine.if affine_set<(d0, d1) : (1 >= 0, 3 >= 0)>(%a, %b) {
     call @external() : () -> ()
   }
   return
 }
 
 // -----
 

 
 // Tests the simplification of a semi-affine expression with a modulo operation with a second operand that simplifies to symbol.
 // CHECK-LABEL: func @semiaffine_unsimplified_symbol
 func @semiaffine_unsimplified_symbol(%arg0: index, %arg1: index) -> index {
   %a = affine.apply affine_map<(d0)[s0] ->(s0 mod (2 * s0 - s0))> (%arg0)[%arg1]
   // CHECK:       %[[CST:.*]] = constant 0
   return %a : index
 }
+
+// -----
+
+// Two external functions that we will use in bodies to avoid DCE.
+func private @external() -> ()
+func private @external1() -> ()
+
+// CHECK-LABEL: func @test_always_true_if_elimination() {
+func @test_always_true_if_elimination() {
+  affine.for %arg0 = 1 to 10 {
+    affine.for %arg1 = 1 to 100 {
+      affine.if affine_set<(d0, d1) : (1 >= 0)> (%arg0, %arg1) {
+        call @external() : () -> ()
+      } else {
+        call @external1() : () -> ()
+      }
+    }
+  }
+  return
+}
+
+// CHECK:      affine.for
+// CHECK-NEXT:   affine.for
+// CHECK-NEXT:     call @external()
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+
+// CHECK-LABEL: func @test_always_false_if_elimination() {
+func @test_always_false_if_elimination() {
+  // CHECK: affine.for
+  affine.for %arg0 = 1 to 10 {
+    // CHECK: affine.for
+    affine.for %arg1 = 1 to 100 {
+      // CHECK: call @external1()
+      // CHECK-NOT: affine.if
+      affine.if affine_set<(d0, d1) : (-1 >= 0)> (%arg0, %arg1) {
+        call @external() : () -> ()
+      } else {
+        call @external1() : () -> ()
+      }
+    }
+  }
+  return
+}
+
+
+// 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.
+// CHECK-LABEL: func @test_num_results_if_elimination
+func @test_num_results_if_elimination() -> f32 {
+  %zero = constant 0.0 : f32
+  // CHECK: affine.if
+  %0 = affine.if affine_set<() : ()> () -> f32 {
+    affine.yield %zero : f32
+    // CHECK: else {
+  } else {
+    affine.yield %zero : f32
+  }
+  return %0 : f32
+}