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:
   ///
   /// Implemented as
   ///
   ///     a mod b =
   ///         let remainder = srem a, b;
   ///             negative = a < 0 in
   ///         select negative, remainder + b, remainder.
   Value visitModExpr(AffineBinaryOpExpr expr) {
-    auto rhsConst = expr.getRHS().dyn_cast<AffineConstantExpr>();
-    if (!rhsConst) {
-      emitError(
-          loc,
-          "semi-affine expressions (modulo by non-const) are not supported");
-      return nullptr;
-    }
+    if (auto rhsConst = expr.getRHS().dyn_cast<AffineConstantExpr>()) {
       if (rhsConst.getValue() <= 0) {
         emitError(loc, "modulo 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 remainder = builder.create<arith::RemSIOp>(loc, lhs, rhs);
     Value zeroCst = builder.create<arith::ConstantIndexOp>(loc, 0);
     Value isRemainderNegative = builder.create<arith::CmpIOp>(
public:
   /// For positive divisors, it can be implemented without branching and with a
   /// 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
+  ///
+  /// Note: this lowering does not use arith.floordivsi because the lowering of
+  /// that to arith.divsi (see populateCeilFloorDivExpandOpsPatterns) generates
+  /// not one but two arith.divsi. That could be changed to one divsi, but one
+  /// way or another, going through arith.floordivsi will result in more complex
+  /// IR because arith.floordivsi is more general than affine floordiv in that
+  /// it supports negative RHS.
   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:
   /// For positive divisors, it can be implemented without branching and with a
   /// 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
+  ///
+  /// Note: not using arith.ceildivsi for the same reason as explained in the
+  /// visitFloorDivExpr comment.
   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.
 // --------------------------------------------------------------------------//
+
+#map_mod = 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)
+  %0 = affine.apply #map_mod (%arg0)
+  return %0 : index
+}
+#map_mod_dynamic_divisor = affine_map<(i)[s] -> (i mod s)>
+// CHECK-LABEL: func @affine_apply_mod_dynamic_divisor
+func.func @affine_apply_mod_dynamic_divisor(%arg0 : index, %arg1 : index) -> (index) {
+// CHECK-NEXT: %[[v0:.*]] = arith.remsi %{{.*}}, %arg1 : index
+// CHECK-NEXT: %[[c0:.*]] = arith.constant 0 : index
+// CHECK-NEXT: %[[v1:.*]] = arith.cmpi slt, %[[v0]], %[[c0]] : index
+// CHECK-NEXT: %[[v2:.*]] = arith.addi %[[v0]], %arg1 : index
+// CHECK-NEXT: %[[v3:.*]] = arith.select %[[v1]], %[[v2]], %[[v0]] : index
+  %0 = affine.apply #map_mod_dynamic_divisor (%arg0)[%arg1]
   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.
 // --------------------------------------------------------------------------//
+#map_floordiv = 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)
+  %0 = affine.apply #map_floordiv (%arg0)
+  return %0 : index
+}
+#map_floordiv_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 #map_floordiv_dynamic_divisor (%arg0)[%arg1]
   return %0 : index
 }
-
-#mapceildiv = affine_map<(i) -> (i ceildiv 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_ceildiv
+// test in the "canonicalize.mlir" test to reflect the expected output of
 // affine.apply lowering.
 // --------------------------------------------------------------------------//
+#map_ceildiv = 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)
+  %0 = affine.apply #map_ceildiv (%arg0)
+  return %0 : index
+}
+#map_ceildiv_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 #map_ceildiv_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>