diff --git a/compiler/src/iree/compiler/Dialect/Util/IR/UtilOps.cpp b/compiler/src/iree/compiler/Dialect/Util/IR/UtilOps.cpp index 9946ecf4a778..410252bab913 100644 --- a/compiler/src/iree/compiler/Dialect/Util/IR/UtilOps.cpp +++ b/compiler/src/iree/compiler/Dialect/Util/IR/UtilOps.cpp @@ -1184,14 +1184,26 @@ AssumeIntOp::getUnionedUnsignedRange(unsigned operandIndex) { umaxCount == assumptions.size() ? umaxUnion : std::nullopt); } +static bool isConstantZero(IntAssumptionAttr assumption) { + std::optional umin = assumption.getUmin(); + std::optional umax = assumption.getUmax(); + if (!umin || !umax) + return false; + return *umin == 0 && *umax == 0; +} + std::optional AssumeIntOp::getUnionedUnsignedDivisor(unsigned operandIndex) { auto assumptions = getOperandAssumptions(operandIndex); std::optional divisorUnion; for (auto assumption : assumptions) { auto divisor = assumption.getUdiv(); - if (!divisor) + if (!divisor) { + // Constant zero is divisible by anything + if (isConstantZero(assumption)) + continue; return std::nullopt; + } if (divisorUnion) divisorUnion = std::gcd(*divisor, *divisorUnion); else diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/test/integer_divisibility.mlir b/compiler/src/iree/compiler/Dialect/Util/Transforms/test/integer_divisibility.mlir index 520fbe9a378d..f5b568e392b6 100644 --- a/compiler/src/iree/compiler/Dialect/Util/Transforms/test/integer_divisibility.mlir +++ b/compiler/src/iree/compiler/Dialect/Util/Transforms/test/integer_divisibility.mlir @@ -100,3 +100,15 @@ util.func @divui_divisibility(%arg0 : index) -> (index, index) { // CHECK: %[[V:.+]] = arith.divui // CHECK: %[[REM:.+]] = arith.remui %[[V]], %[[C32]] // CHECK: return %[[C0]], %[[REM]] + +// ----- + +util.func @everything_divides_zero(%arg0 : index) -> index { + %c32 = arith.constant 32 : index + %0 = util.assume.int %arg0[, ] : index + %1 = arith.remui %0, %c32 : index + util.return %1 : index +} +// CHECK-LABEL: @everything_divides_zero +// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index +// CHECK: return %[[C0]]