[MLIR] Add affine.if canonicalization to compose in affine.apply ops

[bondhugula]

Summary

Add affine.if canonicalization to compose affine.apply ops into its set
and operands. This eliminates affine.apply ops feeding into affine.if
ops.

[dcaballe]

Thanks!

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

2495–2496

missing `

[bondhugula]

Fix typo.

[bondhugula]

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

2495–2496

Thanks.

mlir/include/mlir/IR/IntegerSet.h

public:
   /// the result.  The returned map always has the same number of results.
   IntegerSet replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
                                    ArrayRef<AffineExpr> symReplacements,
                                    unsigned numResultDims,
                                    unsigned numResultSyms);
 
   explicit operator bool() { return set; }
   bool operator==(IntegerSet other) const { return set == other.set; }
+  bool operator!=(IntegerSet other) const { return set != other.set; }
 
   unsigned getNumDims() const;
   unsigned getNumSymbols() const;
   unsigned getNumInputs() const;
   unsigned getNumConstraints() const;
   unsigned getNumEqualities() const;
   unsigned getNumInequalities() const;
 

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

 }
 
 void AffineIfOp::build(OpBuilder &builder, OperationState &result,
                        IntegerSet set, ValueRange args, bool withElseRegion) {
   AffineIfOp::build(builder, result, /*resultTypes=*/{}, set, args,
                     withElseRegion);
 }
 
+/// Compose any affine.apply ops feeding into `operands` of the integer set
+/// `set` by composing the maps of such affine.apply ops with the integer
+/// set constraints.
+static void composeSetAndOperands(IntegerSet &set,
+                                  SmallVectorImpl<Value> &operands) {
+  // We will simply reuse the API of the map composition by viewing the LHSs of
+  // the equalities and inequalities of `set` as the affine exprs of an affine
+  // map. Convert to equivalent map, compose, and convert back to set.
+  auto map = AffineMap::get(set.getNumDims(), set.getNumSymbols(),
+                            set.getConstraints(), set.getContext());
+  // Check if any composition is possible.
+  if (llvm::none_of(operands,
+                    [](Value v) { return v.getDefiningOp<AffineApplyOp>(); }))
+    return;
+
+  composeAffineMapAndOperands(&map, &operands);
+  set = IntegerSet::get(map.getNumDims(), map.getNumSymbols(), map.getResults(),
+                        set.getEqFlags());
+}
+
 /// Canonicalize an affine if op's conditional (integer set + operands).
 LogicalResult AffineIfOp::fold(ArrayRef<Attribute>,
                                SmallVectorImpl<OpFoldResult> &) {
   auto set = getIntegerSet();
   SmallVector<Value, 4> operands(getOperands());
+  composeSetAndOperands(set, operands);
   canonicalizeSetAndOperands(&set, &operands);
 
-  // Any canonicalization change always leads to either a reduction in the
-  // number of operands or a change in the number of symbolic operands
-  // (promotion of dims to symbols).
-  if (operands.size() < getIntegerSet().getNumInputs() ||
-      set.getNumSymbols() > getIntegerSet().getNumSymbols()) {
-    setConditional(set, operands);
-    return success();
-  }
+  // Check if the canonicalization or composition led to any change.
+  if (getIntegerSet() == set && llvm::equal(operands, getOperands()))
+    return failure();
+
 
-  return failure();
-}
-
 void AffineIfOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
   results.add<SimplifyDeadElse, AlwaysTrueOrFalseIf>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // AffineLoadOp
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/Affine/canonicalize.mlir

affine.for %j = 0 to %N {
       "bar"() : () -> ()
     }
   }
   return
 }
 
 // -----
 
+// CHECK-DAG: #[[$SET:.*]] = affine_set<(d0, d1)[s0] : (d0 - 1 >= 0, d1 - 1 == 0, -d0 + s0 + 10 >= 0)>
+
+// CHECK-LABEL: func @canonicalize_affine_if_compose_apply
+// CHECK-SAME:   %[[N:.*]]: index
+func.func @canonicalize_affine_if_compose_apply(%N: index) {
+  %M = affine.apply affine_map<()[s0] -> (s0 + 10)> ()[%N]
+  // CHECK-NEXT: affine.for %[[I:.*]] =
+  affine.for %i = 0 to 1024 {
+    // CHECK-NEXT: affine.for %[[J:.*]] =
+    affine.for %j = 0 to 100 {
+      %j_ = affine.apply affine_map<(d0)[] -> (d0 + 1)> (%j)
+      // CHECK-NEXT: affine.if #[[$SET]](%[[I]], %[[J]])[%[[N]]]
+      affine.if affine_set<(d0, d1)[s0] : (d0 - 1 >= 0, d1 - 2 == 0, -d0 + s0 >= 0)>(%i, %j_)[%M] {
+        "test.foo"() : ()->()
+      }
+    }
+  }
+  return
+}
+
+// -----
+
 // CHECK-DAG: #[[$LBMAP:.*]] = affine_map<()[s0] -> (0, s0)>
 // CHECK-DAG: #[[$UBMAP:.*]] = affine_map<()[s0] -> (1024, s0 * 2)>
 
 // CHECK-LABEL: func @canonicalize_bounds
 // CHECK-SAME: %[[M:.*]]: index,
 // CHECK-SAME: %[[N:.*]]: index)
 func.func @canonicalize_bounds(%M : index, %N : index) {
   %c0 = arith.constant 0 : index