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 @@ -111,18 +111,12 @@ /// 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; + if (auto rhsConst = expr.getRHS().dyn_cast()) { + 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"); @@ -153,15 +147,11 @@ /// 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; + if (auto rhsConst = expr.getRHS().dyn_cast()) { + 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()); 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 @@ -493,13 +493,13 @@ // 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 @@ -512,13 +512,12 @@ 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 @@ -533,14 +532,27 @@ %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 @@ -557,6 +569,22 @@ %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) {