diff --git a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp --- a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp +++ b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp @@ -15,6 +15,7 @@ #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/Utils.h" +#include "mlir/Dialect/Arith/Transforms/Passes.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Vector/IR/VectorOps.h" @@ -535,6 +536,11 @@ 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( diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp --- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp @@ -101,88 +101,17 @@ } /// 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(); - 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(loc, 0); - Value noneCst = builder.create(loc, -1); - Value negative = builder.create( - loc, arith::CmpIPredicate::slt, lhs, zeroCst); - Value negatedDecremented = builder.create(loc, noneCst, lhs); - Value dividend = - builder.create(loc, negative, negatedDecremented, lhs); - Value quotient = builder.create(loc, dividend, rhs); - Value correctedQuotient = - builder.create(loc, noneCst, quotient); - Value result = builder.create(loc, negative, - correctedQuotient, quotient); - return result; + return builder.create(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(); - 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(loc, 0); - Value oneCst = builder.create(loc, 1); - Value nonPositive = builder.create( - loc, arith::CmpIPredicate::sle, lhs, zeroCst); - Value negated = builder.create(loc, zeroCst, lhs); - Value decremented = builder.create(loc, lhs, oneCst); - Value dividend = - builder.create(loc, nonPositive, negated, decremented); - Value quotient = builder.create(loc, dividend, rhs); - Value negatedQuotient = - builder.create(loc, zeroCst, quotient); - Value incrementedQuotient = - builder.create(loc, quotient, oneCst); - Value result = builder.create( - loc, nonPositive, negatedQuotient, incrementedQuotient); - return result; + return builder.create(loc, lhs, rhs); } Value visitConstantExpr(AffineConstantExpr expr) { diff --git a/mlir/test/Conversion/AffineToStandard/lower-affine.mlir b/mlir/test/Conversion/AffineToStandard/lower-affine.mlir --- a/mlir/test/Conversion/AffineToStandard/lower-affine.mlir +++ b/mlir/test/Conversion/AffineToStandard/lower-affine.mlir @@ -497,7 +497,7 @@ // --------------------------------------------------------------------------// // 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 @@ -514,46 +514,22 @@ #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 } diff --git a/mlir/test/Transforms/canonicalize.mlir b/mlir/test/Transforms/canonicalize.mlir --- a/mlir/test/Transforms/canonicalize.mlir +++ b/mlir/test/Transforms/canonicalize.mlir @@ -609,76 +609,6 @@ 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) -> memref<2xi32> {