Support non-constant divisors in affine mod, floordiv, ceildiv lowerings.

[Benoit]

Summary

The requirement that divisor>0 is not enforced here outside of the
constant case, but how to enforce it? If I understand correctly, it is
UB and while it is nice to be able to deterministically intercept UB,
that isn't always feasible. Hopefully, keeping the existing
enforcement in the constant case is enough.

Earlier in https://reviews.llvm.org/D140043 I thought "hey why is the lowering of affine floordiv missing out on opportunity to generate arith.floordivsi". Now I understand why --- arith.floordivsi being more general (supports RHS<0) and expanding in populateCeilFloorDivExpandOpsPatterns to 2 arith.divsi is a solid reason to keep a custom lowering of affine floordiv directly to arith.divsi.

[dcaballe]

The changes in general LGTM but I'm worried that some affine analyses/transformations may rely on the assumption that the affine maps are not semi-affine and something else has to be protected/extended. I think having some feedback from @bondhugula in this regard would be great.

Earlier in https://reviews.llvm.org/D140043 I thought "hey why is the lowering of affine floordiv missing out on opportunity to generate arith.floordivsi". Now I understand why --- arith.floordivsi being more general (supports RHS<0) and expanding in populateCeilFloorDivExpandOpsPatterns to 2 arith.divsi is a solid reason to keep a custom lowering of affine floordiv directly to arith.divsi.

If would be great if you could add a quick comment about this rationale to the logic that you modified in D140043 so that we don't try to do the same again in the future :)

[Benoit]

Added mod

[Benoit]

added comment

[Benoit]

In D140079#3998299, @dcaballe wrote:

If would be great if you could add a quick comment about this rationale to the logic that you modified in D140043 so that we don't try to do the same again in the future :)

Done.

I've also added the same for mod, as I have needed that one too.

To be clear, I don't believe that this Review is really extending semantics -- semi-affine were in an intermediate state where they were technically legal, just this lowering was generating an error on them... With this diff, I have some end-to-end workloads running happily, so this is useful to me at least.

[dcaballe]

I had a quick look and some utilities are actually checking for this condition as well so I think we are good.

mlir/lib/Dialect/Affine/Utils/Utils.cpp

public:
   /// single division operation as
   ///
   ///        a floordiv b =
   ///            let negative = a < 0 in
   ///            let absolute = negative ? -a - 1 : a in
   ///            let quotient = absolute / b in
   ///                negative ? -quotient - 1 : quotient
   Value visitFloorDivExpr(AffineBinaryOpExpr expr) {
-    auto rhsConst = expr.getRHS().dyn_cast<AffineConstantExpr>();
-    if (!rhsConst) {
-      emitError(
-          loc,
-          "semi-affine expressions (division by non-const) are not supported");
-      return nullptr;
-    }
+    if (auto rhsConst = expr.getRHS().dyn_cast<AffineConstantExpr>()) {
       if (rhsConst.getValue() <= 0) {
         emitError(loc, "division by non-positive value is not supported");
         return nullptr;
       }
-
+    }
     auto lhs = visit(expr.getLHS());
     auto rhs = visit(expr.getRHS());
     assert(lhs && rhs && "unexpected affine expr lowering failure");
 
     Value zeroCst = builder.create<arith::ConstantIndexOp>(loc, 0);
     Value noneCst = builder.create<arith::ConstantIndexOp>(loc, -1);
     Value negative = builder.create<arith::CmpIOp>(
         loc, arith::CmpIPredicate::slt, lhs, zeroCst);
public:
   /// single division operation as
   ///
   ///     a ceildiv b =
   ///         let negative = a <= 0 in
   ///         let absolute = negative ? -a : a - 1 in
   ///         let quotient = absolute / b in
   ///             negative ? -quotient : quotient + 1
   Value visitCeilDivExpr(AffineBinaryOpExpr expr) {
-    auto rhsConst = expr.getRHS().dyn_cast<AffineConstantExpr>();
-    if (!rhsConst) {
-      emitError(loc) << "semi-affine expressions (division by non-const) are "
-                        "not supported";
-      return nullptr;
-    }
+    if (auto rhsConst = expr.getRHS().dyn_cast<AffineConstantExpr>()) {
       if (rhsConst.getValue() <= 0) {
         emitError(loc, "division by non-positive value is not supported");
         return nullptr;
       }
+    }
     auto lhs = visit(expr.getLHS());
     auto rhs = visit(expr.getRHS());
     assert(lhs && rhs && "unexpected affine expr lowering failure");
 
     Value zeroCst = builder.create<arith::ConstantIndexOp>(loc, 0);
     Value oneCst = builder.create<arith::ConstantIndexOp>(loc, 1);
     Value nonPositive = builder.create<arith::CmpIOp>(
         loc, arith::CmpIPredicate::sle, lhs, zeroCst);

mlir/test/Conversion/AffineToStandard/lower-affine.mlir

 
 //===---------------------------------------------------------------------===//
 // Test lowering of Euclidean (floor) division, ceil division and modulo
 // operation used in affine expressions.  In addition to testing the
 // operation-level output, check that the obtained results are correct by
 // applying constant folding transformation after affine lowering.
 //===---------------------------------------------------------------------===//
 
-#mapmod = affine_map<(i) -> (i mod 42)>
-
 // --------------------------------------------------------------------------//
 // IMPORTANT NOTE: if you change this test, also change the @lowered_affine_mod
-// test in the "constant-fold.mlir" test to reflect the expected output of
+// test in the "canonicalize.mlir" test to reflect the expected output of
 // affine.apply lowering.
 // --------------------------------------------------------------------------//
+
+#mapmod = affine_map<(i) -> (i mod 42)>
 // CHECK-LABEL: func @affine_apply_mod
 func.func @affine_apply_mod(%arg0 : index) -> (index) {
 // CHECK-NEXT: %[[c42:.*]] = arith.constant 42 : index
 // CHECK-NEXT: %[[v0:.*]] = arith.remsi %{{.*}}, %[[c42]] : index
 // CHECK-NEXT: %[[c0:.*]] = arith.constant 0 : index
 // CHECK-NEXT: %[[v1:.*]] = arith.cmpi slt, %[[v0]], %[[c0]] : index
 // CHECK-NEXT: %[[v2:.*]] = arith.addi %[[v0]], %[[c42]] : index
 // CHECK-NEXT: %[[v3:.*]] = arith.select %[[v1]], %[[v2]], %[[v0]] : index
   %0 = affine.apply #mapmod (%arg0)
   return %0 : index
 }
 
-#mapfloordiv = affine_map<(i) -> (i floordiv 42)>
-
 // --------------------------------------------------------------------------//
-// IMPORTANT NOTE: if you change this test, also change the @lowered_affine_mod
-// test in the "constant-fold.mlir" test to reflect the expected output of
+// IMPORTANT NOTE: if you change this test, also change the @lowered_affine_floordiv
+// test in the "canonicalize.mlir" test to reflect the expected output of
 // affine.apply lowering.
 // --------------------------------------------------------------------------//
+#mapfloordiv = affine_map<(i) -> (i floordiv 42)>
 // CHECK-LABEL: func @affine_apply_floordiv
 func.func @affine_apply_floordiv(%arg0 : index) -> (index) {
 // CHECK-NEXT: %[[c42:.*]] = arith.constant 42 : index
 // CHECK-NEXT: %[[c0:.*]] = arith.constant 0 : index
 // CHECK-NEXT: %[[cm1:.*]] = arith.constant -1 : index
 // CHECK-NEXT: %[[v0:.*]] = arith.cmpi slt, %{{.*}}, %[[c0]] : index
 // CHECK-NEXT: %[[v1:.*]] = arith.subi %[[cm1]], %{{.*}} : index
 // CHECK-NEXT: %[[v2:.*]] = arith.select %[[v0]], %[[v1]], %{{.*}} : index
 // CHECK-NEXT: %[[v3:.*]] = arith.divsi %[[v2]], %[[c42]] : index
 // CHECK-NEXT: %[[v4:.*]] = arith.subi %[[cm1]], %[[v3]] : index
 // CHECK-NEXT: %[[v5:.*]] = arith.select %[[v0]], %[[v4]], %[[v3]] : index
   %0 = affine.apply #mapfloordiv (%arg0)
   return %0 : index
 }
-
-#mapceildiv = affine_map<(i) -> (i ceildiv 42)>
+#mapfloordiv_dynamic_divisor = affine_map<(i)[s] -> (i floordiv s)>
+// CHECK-LABEL: func @affine_apply_floordiv_dynamic_divisor
+func.func @affine_apply_floordiv_dynamic_divisor(%arg0 : index, %arg1 : index) -> (index) {
+// CHECK-NEXT: %[[c0:.*]] = arith.constant 0 : index
+// CHECK-NEXT: %[[cm1:.*]] = arith.constant -1 : index
+// CHECK-NEXT: %[[v0:.*]] = arith.cmpi slt, %{{.*}}, %[[c0]] : index
+// CHECK-NEXT: %[[v1:.*]] = arith.subi %[[cm1]], %{{.*}} : index
+// CHECK-NEXT: %[[v2:.*]] = arith.select %[[v0]], %[[v1]], %{{.*}} : index
+// CHECK-NEXT: %[[v3:.*]] = arith.divsi %[[v2]], %arg1 : index
+// CHECK-NEXT: %[[v4:.*]] = arith.subi %[[cm1]], %[[v3]] : index
+// CHECK-NEXT: %[[v5:.*]] = arith.select %[[v0]], %[[v4]], %[[v3]] : index
+  %0 = affine.apply #mapfloordiv_dynamic_divisor (%arg0)[%arg1]
+  return %0 : index
+}
 
 // --------------------------------------------------------------------------//
-// IMPORTANT NOTE: if you change this test, also change the @lowered_affine_mod
-// test in the "constant-fold.mlir" test to reflect the expected output of
+// IMPORTANT NOTE: if you change this test, also change the @lowered_affine_ceildiv
+// test in the "canonicalize.mlir" test to reflect the expected output of
 // affine.apply lowering.
 // --------------------------------------------------------------------------//
+#mapceildiv = affine_map<(i) -> (i ceildiv 42)>
 // CHECK-LABEL: func @affine_apply_ceildiv
 func.func @affine_apply_ceildiv(%arg0 : index) -> (index) {
 // CHECK-NEXT:  %[[c42:.*]] = arith.constant 42 : index
 // CHECK-NEXT:  %[[c0:.*]] = arith.constant 0 : index
 // CHECK-NEXT:  %[[c1:.*]] = arith.constant 1 : index
 // CHECK-NEXT:  %[[v0:.*]] = arith.cmpi sle, %{{.*}}, %[[c0]] : index
 // CHECK-NEXT:  %[[v1:.*]] = arith.subi %[[c0]], %{{.*}} : index
 // CHECK-NEXT:  %[[v2:.*]] = arith.subi %{{.*}}, %[[c1]] : index
 // CHECK-NEXT:  %[[v3:.*]] = arith.select %[[v0]], %[[v1]], %[[v2]] : index
 // CHECK-NEXT:  %[[v4:.*]] = arith.divsi %[[v3]], %[[c42]] : index
 // CHECK-NEXT:  %[[v5:.*]] = arith.subi %[[c0]], %[[v4]] : index
 // CHECK-NEXT:  %[[v6:.*]] = arith.addi %[[v4]], %[[c1]] : index
 // CHECK-NEXT:  %[[v7:.*]] = arith.select %[[v0]], %[[v5]], %[[v6]] : index
   %0 = affine.apply #mapceildiv (%arg0)
   return %0 : index
 }
+#mapceildiv_dynamic_divisor = affine_map<(i)[s] -> (i ceildiv s)>
+// CHECK-LABEL: func @affine_apply_ceildiv_dynamic_divisor
+func.func @affine_apply_ceildiv_dynamic_divisor(%arg0 : index, %arg1 : index) -> (index) {
+// CHECK-NEXT:  %[[c0:.*]] = arith.constant 0 : index
+// CHECK-NEXT:  %[[c1:.*]] = arith.constant 1 : index
+// CHECK-NEXT:  %[[v0:.*]] = arith.cmpi sle, %{{.*}}, %[[c0]] : index
+// CHECK-NEXT:  %[[v1:.*]] = arith.subi %[[c0]], %{{.*}} : index
+// CHECK-NEXT:  %[[v2:.*]] = arith.subi %{{.*}}, %[[c1]] : index
+// CHECK-NEXT:  %[[v3:.*]] = arith.select %[[v0]], %[[v1]], %[[v2]] : index
+// CHECK-NEXT:  %[[v4:.*]] = arith.divsi %[[v3]], %arg1 : index
+// CHECK-NEXT:  %[[v5:.*]] = arith.subi %[[c0]], %[[v4]] : index
+// CHECK-NEXT:  %[[v6:.*]] = arith.addi %[[v4]], %[[c1]] : index
+// CHECK-NEXT:  %[[v7:.*]] = arith.select %[[v0]], %[[v5]], %[[v6]] : index
+  %0 = affine.apply #mapceildiv_dynamic_divisor (%arg0)[%arg1]
+  return %0 : index
+}
 
 // CHECK-LABEL: func @affine_load
 func.func @affine_load(%arg0 : index) {
   %0 = memref.alloc() : memref<10xf32>
   affine.for %i0 = 0 to 10 {
     %1 = affine.load %0[%i0 + symbol(%arg0) + 7] : memref<10xf32>
   }
 // CHECK:       %[[a:.*]] = arith.addi %{{.*}}, %{{.*}} : index