[Mlir-commits] [mlir] eec575e - Allow non-constant divisors in affine mod, floordiv, ceildiv.
Benoit Jacob
llvmlistbot at llvm.org
Fri Dec 16 18:27:16 PST 2022
Author: Benoit Jacob
Date: 2022-12-17T02:24:02Z
New Revision: eec575e548d114d96acb673ccb9b8a1ef795465b
URL: https://github.com/llvm/llvm-project/commit/eec575e548d114d96acb673ccb9b8a1ef795465b
DIFF: https://github.com/llvm/llvm-project/commit/eec575e548d114d96acb673ccb9b8a1ef795465b.diff
LOG: Allow non-constant divisors in affine mod, floordiv, ceildiv.
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.
Differential Revision: https://reviews.llvm.org/D140079
Added:
Modified:
mlir/lib/Dialect/Affine/Utils/Utils.cpp
mlir/test/Conversion/AffineToStandard/lower-affine.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index e93455a5b066d..eb7224fbdc985 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -73,16 +73,11 @@ class AffineApplyExpander
/// 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 (rhsConst.getValue() <= 0) {
- emitError(loc, "modulo by non-positive value is 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());
@@ -110,19 +105,20 @@ class AffineApplyExpander
/// 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 (rhsConst.getValue() <= 0) {
- emitError(loc, "division by non-positive value is 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");
@@ -152,16 +148,15 @@ class AffineApplyExpander
/// 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 (rhsConst.getValue() <= 0) {
- emitError(loc, "division by non-positive value is 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());
diff --git a/mlir/test/Conversion/AffineToStandard/lower-affine.mlir b/mlir/test/Conversion/AffineToStandard/lower-affine.mlir
index 5fc8d010bcd30..6158de33e4aef 100644
--- a/mlir/test/Conversion/AffineToStandard/lower-affine.mlir
+++ b/mlir/test/Conversion/AffineToStandard/lower-affine.mlir
@@ -493,13 +493,13 @@ func.func @args_ret_affine_apply(index, index) -> (index, index) {
// 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
@@ -508,17 +508,27 @@ func.func @affine_apply_mod(%arg0 : index) -> (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
@@ -530,17 +540,30 @@ func.func @affine_apply_floordiv(%arg0 : index) -> (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
@@ -554,7 +577,23 @@ func.func @affine_apply_ceildiv(%arg0 : index) -> (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
}
More information about the Mlir-commits
mailing list