Affine expr simplification for add of const multiple of same expression

[bondhugula]

Summary

• Detect "c_1 * expr + c_2 * expr" as (c_1 + c_2) * expr
• subsumes things like 'expr - expr' and "expr * -1 + expr" as 0.
• add default null init ctor for AffineConstantExpr

Signed-off-by: Uday Bondhugula <[e-mail]>

[bondhugula]

Remove debug code

[nicolasvasilache]

LGTM

See: https://github.com/llvm/llvm-project/blob/master/mlir/lib/Dialect/SDBM/SDBMExpr.cpp#L31 in case you like something from there to write such expressions more idiomatically and wish to lift it up.
I am not claiming it is in a final state but you might find it useful.

[rriddle]

mlir/include/mlir/IR/AffineExpr.h

191

Can you just add = nullptr to the other constructor?

[bondhugula]

Thanks for the review.

mlir/include/mlir/IR/AffineExpr.h

191

Sure, done.

[bondhugula]

Address review comment - AffineConstantExpr ctor

[bondhugula]

In D76233#1924892, @nicolasvasilache wrote:

LGTM

See: https://github.com/llvm/llvm-project/blob/master/mlir/lib/Dialect/SDBM/SDBMExpr.cpp#L31 in case you like something from there to write such expressions more idiomatically and wish to lift it up.
I am not claiming it is in a final state but you might find it useful.

Looks good and useful. You may also want to add the operator '/' for floordiv. This would be in general useful to implement some detection at least from under Analysis/. You could just consider moving it to lib/Analysis/Utils so that it's not a dependence on the SDBM dialect. One use case could have been to simplify (x - x floordiv d) as x mod d, but we have this by construction in lib/IR/AffineExpr.cpp. Also, it may be useful to have a way to constrain an AffineExprMatcher to positive constants for example. But this is all meant for "post construction" matching as opposed to simplifying before / by construction itself, right?

mlir/include/mlir/IR/AffineExpr.h

public:
   /* implicit */ AffineSymbolExpr(AffineExpr::ImplType *ptr);
   unsigned getPosition() const;
 };
 
 /// An integer constant appearing in affine expression.
 class AffineConstantExpr : public AffineExpr {
 public:
   using ImplType = detail::AffineConstantExprStorage;
-  /* implicit */ AffineConstantExpr(AffineExpr::ImplType *ptr);
+  /* implicit */ AffineConstantExpr(AffineExpr::ImplType *ptr = nullptr);
   int64_t getValue() const;
 };
 
 /// Make AffineExpr hashable.
 inline ::llvm::hash_code hash_value(AffineExpr arg) {
   return ::llvm::hash_value(arg.expr);
 }
 
 inline AffineExpr operator+(int64_t val, AffineExpr expr) { return expr + val; }

mlir/lib/IR/AffineExpr.cpp

static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs) {
   }
   // Fold successive additions like (d0 + 2) + 3 into d0 + 5.
   auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>();
   if (lBin && rhsConst && lBin.getKind() == AffineExprKind::Add) {
     if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>())
       return lBin.getLHS() + (lrhs.getValue() + rhsConst.getValue());
   }
 
+  // Detect "c1 * expr + c_2 * expr" as "(c1 + c2) * expr".
+  // c1 is rRhsConst, c2 is rLhsConst; firstExpr, secondExpr are their
+  // respective multiplicands.
+  Optional<int64_t> rLhsConst, rRhsConst;
+  AffineExpr firstExpr, secondExpr;
+  AffineConstantExpr rLhsConstExpr;
+  auto lBinOpExpr = lhs.dyn_cast<AffineBinaryOpExpr>();
+  if (lBinOpExpr && lBinOpExpr.getKind() == AffineExprKind::Mul &&
+      (rLhsConstExpr = lBinOpExpr.getRHS().dyn_cast<AffineConstantExpr>())) {
+    rLhsConst = rLhsConstExpr.getValue();
+    firstExpr = lBinOpExpr.getLHS();
+  } else {
+    rLhsConst = 1;
+    firstExpr = lhs;
+  }
+
+  auto rBinOpExpr = rhs.dyn_cast<AffineBinaryOpExpr>();
+  AffineConstantExpr rRhsConstExpr;
+  if (rBinOpExpr && rBinOpExpr.getKind() == AffineExprKind::Mul &&
+      (rRhsConstExpr = rBinOpExpr.getRHS().dyn_cast<AffineConstantExpr>())) {
+    rRhsConst = rRhsConstExpr.getValue();
+    secondExpr = rBinOpExpr.getLHS();
+  } else {
+    rRhsConst = 1;
+    secondExpr = rhs;
+  }
+
+  if (rLhsConst && rRhsConst && firstExpr == secondExpr)
+    return getAffineBinaryOpExpr(
+        AffineExprKind::Mul, firstExpr,
+        getAffineConstantExpr(rLhsConst.getValue() + rRhsConst.getValue(),
+                              lhs.getContext()));
+
   // When doing successive additions, bring constant to the right: turn (d0 + 2)
   // + d1 into (d0 + d1) + 2.
   if (lBin && lBin.getKind() == AffineExprKind::Add) {
     if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) {
       return lBin.getLHS() + rhs + lrhs;
     }
   }
 
   // Detect and transform "expr - c * (expr floordiv c)" to "expr mod c". This
   // leads to a much more efficient form when 'c' is a power of two, and in
   // general a more compact and readable form.
 
   // Process '(expr floordiv c) * (-c)'.
-  AffineBinaryOpExpr rBinOpExpr = rhs.dyn_cast<AffineBinaryOpExpr>();
   if (!rBinOpExpr)
     return nullptr;
 
   auto lrhs = rBinOpExpr.getLHS();
   auto rrhs = rBinOpExpr.getRHS();
 
   // Process lrhs, which is 'expr floordiv c'.
   AffineBinaryOpExpr lrBinOpExpr = lrhs.dyn_cast<AffineBinaryOpExpr>();

mlir/test/Dialect/AffineOps/canonicalize.mlir

affine.for %i = 0 to 1024 {
     }
   }
   return
 }
 
 // -----
 
 // CHECK-DAG: [[LBMAP:#map[0-9]+]] = affine_map<()[s0] -> (0, s0)>
-// CHECK-DAG: [[UBMAP:#map[0-9]+]] = affine_map<()[s0] -> (1024, s0 + s0)>
+// CHECK-DAG: [[UBMAP:#map[0-9]+]] = affine_map<()[s0] -> (1024, s0 * 2)>
 
 // CHECK-LABEL: func @canonicalize_bounds
 // CHECK-SAME: [[M:%.*]]: index,
 // CHECK-SAME: [[N:%.*]]: index)
 func @canonicalize_bounds(%M : index, %N : index) {
   %c0 = constant 0 : index
   %c1024 = constant 1024 : index
   // Drop unused operand %N, drop duplicate operand %M, propagate %c1024, and

mlir/test/IR/affine-map.mlir

 #map3h = affine_map<(i, j) -> (i + 1, 2*j*2 + 2)>
 #map3i = affine_map<(i, j) -> (i + 1, j*2*2 + 2)>
 #map3j = affine_map<(i, j) -> (i + 1, j*1*4 + 2)>
 #map3k = affine_map<(i, j) -> (i + 1, j*4*1 + 2)>
 
 // The following reduction should be unique'd out too but such expression
 // simplification is not performed for IR parsing, but only through analyses
 // and transforms.
-// CHECK: #map{{[0-9]+}} = affine_map<(d0, d1) -> (d1 - d0 + (d0 - d1 + 1) * 2 + d1 - 1, d1 + d1 + d1 + d1 + 2)>
+// CHECK: #map{{[0-9]+}} = affine_map<(d0, d1) -> (d1 - d0 + (d0 - d1 + 1) * 2 + d1 - 1, d1 * 4 + 2)>
 #map3l = affine_map<(i, j) -> ((j - i) + 2*(i - j + 1) + j - 1 + 0, j + j + 1 + j + j + 1)>
 
 // CHECK: #map{{[0-9]+}} = affine_map<(d0, d1) -> (d0 + 2, d1)>
 #map4  = affine_map<(i, j) -> (3+3-2*2+i, j)>
 
 // CHECK: #map{{[0-9]+}} = affine_map<(d0, d1)[s0] -> (d0 + s0, d1)>
 #map5 = affine_map<(i, j)[s0] -> (i + s0, j)>
 

 #map54 = affine_map<(d0)[s0] -> (-s0 floordiv 4, d0 floordiv -1)>
 
 // CHECK: #map{{[0-9]+}} = affine_map<() -> ()>
 #map55 = affine_map<() -> ()>
 
 // CHECK: #map{{[0-9]+}} = affine_map<(d0, d1) -> (d0, d0 * 2 + d1 * 4 + 2, 1, 2, (d0 * 4) mod 8)>
 #map56 = affine_map<(d0, d1) -> ((4*d0 + 2) floordiv 4, (4*d0 + 8*d1 + 5) floordiv 2, (2*d0 + 4*d1 + 3) mod 2, (3*d0 - 4) mod 3, (4*d0 + 8*d1) mod 8)>
 
+// CHECK: #map{{[0-9]+}} = affine_map<(d0, d1) -> (d1, d0, 0)>
+#map57 = affine_map<(d0, d1) -> (d0 - d0 + d1, -d0 + d0 + d0, (1 + d0 + d1 floordiv 4) - (d0 + d1 floordiv 4 + 1))>
+
+// CHECK: #map{{[0-9]+}} = affine_map<(d0, d1) -> (d0 * 3, (d0 + d1) * 2, d0 mod 2)>
+#map58 = affine_map<(d0, d1) -> (4*d0 - 2*d0 + d0, (d0 + d1) + (d0 + d1), 2 * (d0 mod 2) - d0 mod 2)>
+
 // Single identity maps are removed.
 // CHECK: func @f0(memref<2x4xi8, 1>)
 func @f0(memref<2x4xi8, #map0, 1>)
 
 // Single identity maps are removed.
 // CHECK: func @f1(memref<2x4xi8, 1>)
 func @f1(memref<2x4xi8, #map1, 1>)
 

 // CHECK: func @f54(memref<10xi32, #map{{[0-9]+}}>)
 func @f54(memref<10xi32, #map54>)
 
 // CHECK: "foo.op"() {map = #map{{[0-9]+}}} : () -> ()
 "foo.op"() {map = #map55} : () -> ()
 
 // CHECK: func @f56(memref<1x1xi8, #map{{[0-9]+}}>)
 func @f56(memref<1x1xi8, #map56>)
+
+// CHECK: "f57"() {map = #map{{[0-9]+}}} : () -> ()
+"f57"() {map = #map57} : () -> ()
+
+// CHECK: "f58"() {map = #map{{[0-9]+}}} : () -> ()
+"f58"() {map = #map58} : () -> ()