Add affine expr simplification for subtract from self

[bondhugula]

Summary

• Detect "expr + expr * (-1)" and "expr * -1 + expr" as 0.

[rriddle]

mlir/lib/IR/AffineExpr.cpp

319

It looks like these detection comments are reversed. I may be missing something, but it looks like this code block computes : expr * -1 + expr = 0.

320

It looks like lBinOpExpr is the same as lBin?

330

Why are these wrapped in scopes?

mlir/test/IR/affine-map.mlir

187

Can we not test the added detection directly: d0 + d0 * -1 and d0 * -1 + d0?

[nicolasvasilache]

mlir/lib/IR/AffineExpr.cpp

317

This looks like it could separate concerns better..
Shouldn't this be a single pattern that adds the constants and lets a 0 * expr canonicalization rewrite it to 0?

[bondhugula]

mlir/lib/IR/AffineExpr.cpp

317

We can actually generalize this to "check for constant multipliers if both are binary op expressions, add the constants, and the 0 times expr is already constructed as 0 if that's the case" - but that would make the logic and casing long and complex because expr * 1 is always expr by construction. So, you'd have to check for both expr and expr * const_not_1 on either side -- four different cases. I kept this simple for now but can extend it this way if you prefer.

[bondhugula]

This has been subsumed by D76233.

mlir/lib/IR/AffineExpr.cpp

317

I've now generalized to simplify c1 * expr + c2 * expr to (c1 + c2) * expr - the logic is more complex, but it's better than having two patches.

319

You are right - the comments were switched. Anyway the commit is now generalized.

330

There was no good reason - perceived readability. Anyway, this is outdated now.

mlir/test/IR/affine-map.mlir

187

d0 - d0 is internally d0 + d0 * -1. So we'll get the same thing.

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 "expr + expr * (-1)" and "expr * -1 + expr" as 0.
+  {
+    // expr + expr * -1 = 0.
+    auto lBinOpExpr = lhs.dyn_cast<AffineBinaryOpExpr>();
+    if (lBinOpExpr && lBinOpExpr.getKind() == AffineExprKind::Mul) {
+      auto llhs = lBinOpExpr.getLHS();
+      auto rlhs = lBinOpExpr.getRHS();
+      auto rlhsConst = rlhs.dyn_cast<AffineConstantExpr>();
+      if (llhs == rhs && rlhsConst && rlhsConst.getValue() == -1)
+        return getAffineConstantExpr(0, lhs.getContext());
+    }
+  }
+  auto rBinOpExpr = rhs.dyn_cast<AffineBinaryOpExpr>();
+  {
+    // expr * -1 + expr = 0.
+    if (rBinOpExpr && rBinOpExpr.getKind() == AffineExprKind::Mul) {
+      auto lrhs = rBinOpExpr.getLHS();
+      auto rrhs = rBinOpExpr.getRHS();
+      auto rrhsConst = rrhs.dyn_cast<AffineConstantExpr>();
+      if (lrhs == lhs && rrhsConst && rrhsConst.getValue() == -1)
+        return getAffineConstantExpr(0, 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/IR/affine-map.mlir

 #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))>
+
 // 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} : () -> ()