Generate `arith.{floordivsi,ceildivsi}` in lowering affine exprs

[Benoit]

Summary

It seems that the existing lowerings must have predated the introduction of arith.{floordivsi,ceildivsi}, requiring a custom implementation based on arith.divsi and several other ops to implement the floor/ceil aspects?

My immediate motivation for this is that the existing code had to check if the RHS was positive, and it could only do so when the RHS was constant, so it required it to be constant, and I have a use case where it's dynamic, so I ran into the error there.

Looking at the tests, I saw the comment pointing to the corresponding canonicalize.mlir testcase (actually the comment mentioned another file name but the testcase moved, so this PR fixes that). But then looking at the canonicalize.mlir testcase, it seems that its existence was justified by the complex IR that used to be produced, and now that it's a single arith op, no canonicalize.mlir testcase is needed anymore. So this PR also removes that.

[Benoit]

Expand floordivsi/ceildivsi to avoid breaking downstream patterns not expecting those
(caught by e2e compilation tests in IREE).

[Benoit]

clang-format

[dcaballe]

mlir/test/Transforms/canonicalize.mlir

616

can we regenerate the IR from affine.apply affine_map<(i) -> (i mod 42)> and make sure that canonicalization is still happening?

[Benoit]

I checked, the checked IR in lower-affine.mlir for affine_apply_mod still matches what canonicalize.mlir checks in lowered_affine_mod .

[dcaballe]

Great! Could you please move the canonicalization tests to test/Dialect/Arith/canonicalize.mlir using the new floor/ceil ops? Other than that, LGTM. Thanks!

mlir/test/Transforms/canonicalize.mlir

616

I guess if we have a similar folding test for the floor/ceil div ops, that should be enough. If you can check if that is the case, that should be all!

[Benoit]

Dropping this Review. I realized that the lowering (by ExpandOps.cpp) of arith.floordivsi to arith.divsi unconditionally generates 2 divisions. The branch-free lowering approach there is misfiring --- it's great to be branch-free, but not if that means 2x more divisions, as these days divisions (with non-constant RHS) are far more expensive than some branching inside some arithmetic. I will instead work on removing my code's reliance on affine floordiv/ceildiv.

[Benoit]

Take 2 in https://reviews.llvm.org/D140079

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/Passes.h"
+#include "mlir/Dialect/Arith/Transforms/Passes.h"
 
 namespace mlir {
 #define GEN_PASS_DEF_CONVERTAFFINETOSTANDARD
 #include "mlir/Conversion/Passes.h.inc"
 } // namespace mlir
 
 using namespace mlir;
 using namespace mlir::vector;
patterns.add<
       AffineMaxLowering,
       AffineParallelLowering,
       AffinePrefetchLowering,
       AffineStoreLowering,
       AffineForLowering,
       AffineIfLowering,
       AffineYieldOpLowering>(patterns.getContext());
   // clang-format on
+
+  // AffineApplyLowering may create arith.floordivsi and arith.ceildivsi ops
+  // which have to be expanded to arith.divsi and other ops as subsequent
+  // lowerings (ArithToLLVM) do not handle them.
+  arith::populateCeilFloorDivExpandOpsPatterns(patterns);
 }
 
 void mlir::populateAffineToVectorConversionPatterns(
     RewritePatternSet &patterns) {
   // clang-format off
   patterns.add<
       AffineVectorLoadLowering,
       AffineVectorStoreLowering>(patterns.getContext());

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

Value visitModExpr(AffineBinaryOpExpr expr) {
     Value correctedRemainder =
         builder.create<arith::AddIOp>(loc, remainder, rhs);
     Value result = builder.create<arith::SelectOp>(
         loc, isRemainderNegative, correctedRemainder, remainder);
     return result;
   }
 
   /// Floor division operation (rounds towards negative infinity).
-  ///
-  /// 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
   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 (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);
-    Value negatedDecremented = builder.create<arith::SubIOp>(loc, noneCst, lhs);
-    Value dividend =
-        builder.create<arith::SelectOp>(loc, negative, negatedDecremented, lhs);
-    Value quotient = builder.create<arith::DivSIOp>(loc, dividend, rhs);
-    Value correctedQuotient =
-        builder.create<arith::SubIOp>(loc, noneCst, quotient);
-    Value result = builder.create<arith::SelectOp>(loc, negative,
-                                                   correctedQuotient, quotient);
-    return result;
+    return builder.create<arith::FloorDivSIOp>(loc, lhs, rhs);
   }
 
   /// Ceiling division operation (rounds towards positive infinity).
-  ///
-  /// 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
   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 (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);
-    Value negated = builder.create<arith::SubIOp>(loc, zeroCst, lhs);
-    Value decremented = builder.create<arith::SubIOp>(loc, lhs, oneCst);
-    Value dividend =
-        builder.create<arith::SelectOp>(loc, nonPositive, negated, decremented);
-    Value quotient = builder.create<arith::DivSIOp>(loc, dividend, rhs);
-    Value negatedQuotient =
-        builder.create<arith::SubIOp>(loc, zeroCst, quotient);
-    Value incrementedQuotient =
-        builder.create<arith::AddIOp>(loc, quotient, oneCst);
-    Value result = builder.create<arith::SelectOp>(
-        loc, nonPositive, negatedQuotient, incrementedQuotient);
-    return result;
+    return builder.create<arith::CeilDivSIOp>(loc, lhs, rhs);
   }
 
   Value visitConstantExpr(AffineConstantExpr expr) {
     auto op = builder.create<arith::ConstantIndexOp>(loc, expr.getValue());
     return op.getResult();
   }
 
   Value visitDimExpr(AffineDimExpr expr) {

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

 // 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.
 // --------------------------------------------------------------------------//
 // 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
-// affine.apply lowering.
-// --------------------------------------------------------------------------//
 // 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
+// CHECK-NEXT: %[[v:.*]] = arith.floordivsi %arg0, %[[c42]] : index 
+// CHECK-NEXT: return %[[v]] : index
   %0 = affine.apply #mapfloordiv (%arg0)
   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
-// affine.apply lowering.
-// --------------------------------------------------------------------------//
 // 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
+// CHECK-NEXT: %[[c42:.*]] = arith.constant 42 : index
+// CHECK-NEXT: %[[v:.*]] = arith.ceildivsi %arg0, %[[c42]] : index 
+// CHECK-NEXT: return %[[v]] : index
   %0 = affine.apply #mapceildiv (%arg0)
   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 {

mlir/test/Transforms/canonicalize.mlir

// CHECK-DAG: {{.*}} = arith.constant 41 : index
   %4 = arith.remsi %c43, %c42_0 : index
   %c0_1 = arith.constant 0 : index
   %5 = arith.cmpi slt, %4, %c0_1 : index
   %6 = arith.addi %4, %c42_0 : index
   %7 = arith.select %5, %6, %4 : index
   return %3, %7 : index, index
 }
 
-//
-// IMPORTANT NOTE: the operations in this test are exactly those produced by
-// lowering affine.apply affine_map<(i) -> (i mod 42)> to standard operations.  Please only
-// change these operations together with the affine lowering pass tests.
-//
-// CHECK-LABEL: func @lowered_affine_floordiv
-func.func @lowered_affine_floordiv() -> (index, index) {
-// CHECK-DAG: %c1 = arith.constant 1 : index
-// CHECK-DAG: %c-2 = arith.constant -2 : index
-  %c-43 = arith.constant -43 : index
-  %c42 = arith.constant 42 : index
-  %c0 = arith.constant 0 : index
-  %c-1 = arith.constant -1 : index
-  %0 = arith.cmpi slt, %c-43, %c0 : index
-  %1 = arith.subi %c-1, %c-43 : index
-  %2 = arith.select %0, %1, %c-43 : index
-  %3 = arith.divsi %2, %c42 : index
-  %4 = arith.subi %c-1, %3 : index
-  %5 = arith.select %0, %4, %3 : index
-  %c43 = arith.constant 43 : index
-  %c42_0 = arith.constant 42 : index
-  %c0_1 = arith.constant 0 : index
-  %c-1_2 = arith.constant -1 : index
-  %6 = arith.cmpi slt, %c43, %c0_1 : index
-  %7 = arith.subi %c-1_2, %c43 : index
-  %8 = arith.select %6, %7, %c43 : index
-  %9 = arith.divsi %8, %c42_0 : index
-  %10 = arith.subi %c-1_2, %9 : index
-  %11 = arith.select %6, %10, %9 : index
-  return %5, %11 : index, index
-}
-
-//
-// IMPORTANT NOTE: the operations in this test are exactly those produced by
-// lowering affine.apply affine_map<(i) -> (i mod 42)> to standard operations.  Please only
-// change these operations together with the affine lowering pass tests.
-//
-// CHECK-LABEL: func @lowered_affine_ceildiv
-func.func @lowered_affine_ceildiv() -> (index, index) {
-// CHECK-DAG:  %c-1 = arith.constant -1 : index
-  %c-43 = arith.constant -43 : index
-  %c42 = arith.constant 42 : index
-  %c0 = arith.constant 0 : index
-  %c1 = arith.constant 1 : index
-  %0 = arith.cmpi sle, %c-43, %c0 : index
-  %1 = arith.subi %c0, %c-43 : index
-  %2 = arith.subi %c-43, %c1 : index
-  %3 = arith.select %0, %1, %2 : index
-  %4 = arith.divsi %3, %c42 : index
-  %5 = arith.subi %c0, %4 : index
-  %6 = arith.addi %4, %c1 : index
-  %7 = arith.select %0, %5, %6 : index
-// CHECK-DAG:  %c2 = arith.constant 2 : index
-  %c43 = arith.constant 43 : index
-  %c42_0 = arith.constant 42 : index
-  %c0_1 = arith.constant 0 : index
-  %c1_2 = arith.constant 1 : index
-  %8 = arith.cmpi sle, %c43, %c0_1 : index
-  %9 = arith.subi %c0_1, %c43 : index
-  %10 = arith.subi %c43, %c1_2 : index
-  %11 = arith.select %8, %9, %10 : index
-  %12 = arith.divsi %11, %c42_0 : index
-  %13 = arith.subi %c0_1, %12 : index
-  %14 = arith.addi %12, %c1_2 : index
-  %15 = arith.select %8, %13, %14 : index
-
-  // CHECK-NEXT: return %c-1, %c2
-  return %7, %15 : index, index
-}
-
 // Checks that NOP casts are removed.
 // CHECK-LABEL: cast_values
 func.func @cast_values(%arg0: memref<?xi32>) -> memref<2xi32> {
   // NOP cast
   %1 = memref.cast %arg0 : memref<?xi32> to memref<?xi32>
   // CHECK-NEXT: %[[RET:.*]] = memref.cast %arg0 : memref<?xi32> to memref<2xi32>
   %3 = memref.cast %1 : memref<?xi32> to memref<2xi32>
   // NOP cast