Model UB in integer division operations in the arith dialect

[sanjoy]

Summary

Before this commit arith.{ceil}div{u|s}i were marked Pure which is
incorrect because these operations invoke UB on certain inputs.

Fixes: https://github.com/llvm/llvm-project/issues/58700

[kuhar]

LGTM. Thanks for taking care of this!

[kuhar]

(BTW, you probably already know this, but if you add Fixes: <github-issue-url> to the commit description landing the change will automatically close the issue.)

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

 include "mlir/Interfaces/InferIntRangeInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/VectorInterfaces.td"
 include "mlir/IR/BuiltinAttributeInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/EnumAttr.td"
 
-// Base class for Arith dialect ops. Ops in this dialect have no side
+// Base class for Arith dialect ops. Ops in this dialect have no memory
 // effects and can be applied element-wise to vectors and tensors.
 class Arith_Op<string mnemonic, list<Trait> traits = []> :
-    Op<Arith_Dialect, mnemonic, traits # [Pure,
-    DeclareOpInterfaceMethods<VectorUnrollOpInterface>] #
+    Op<Arith_Dialect, mnemonic,
+       traits #
+       [DeclareOpInterfaceMethods<VectorUnrollOpInterface>, NoMemoryEffect] #
        ElementwiseMappable.traits>;
 
 // Base class for integer and floating point arithmetic ops. All ops have one
 // result, require operands and results to be of the same type, and can accept
 // tensors or vectors of integers or floats.
 class Arith_ArithOp<string mnemonic, list<Trait> traits = []> :
     Arith_Op<mnemonic, traits # [SameOperandsAndResultType]>;
 
 // Base class for unary arithmetic operations.
 class Arith_UnaryOp<string mnemonic, list<Trait> traits = []> :
-    Arith_ArithOp<mnemonic, traits> {
+    Arith_ArithOp<mnemonic, traits # [Pure]> {
   let assemblyFormat = "$operand attr-dict `:` type($result)";
 }
 
 // Base class for binary arithmetic operations.
 class Arith_BinaryOp<string mnemonic, list<Trait> traits = []> :
     Arith_ArithOp<mnemonic, traits> {
   let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($result)";
 }
 
 // Base class for ternary arithmetic operations.
 class Arith_TernaryOp<string mnemonic, list<Trait> traits = []> :
-    Arith_ArithOp<mnemonic, traits> {
+    Arith_ArithOp<mnemonic, traits # [Pure]> {
   let assemblyFormat = "$a `,` $b `,` $c attr-dict `:` type($result)";
 }
 
 // Base class for integer binary operations.
 class Arith_IntBinaryOp<string mnemonic, list<Trait> traits = []> :
     Arith_BinaryOp<mnemonic, traits #
       [DeclareOpInterfaceMethods<InferIntRangeInterface>]>,
     Arguments<(ins SignlessIntegerLike:$lhs, SignlessIntegerLike:$rhs)>,
     Results<(outs SignlessIntegerLike:$result)>;
 
+// Base class for integer binary operations without undefined behavior.
+class Arith_TotalIntBinaryOp<string mnemonic, list<Trait> traits = []> :
+    Arith_IntBinaryOp<mnemonic, traits # [Pure]>;
+
 // Base class for floating point unary operations.
 class Arith_FloatUnaryOp<string mnemonic, list<Trait> traits = []> :
     Arith_UnaryOp<mnemonic,
       !listconcat([DeclareOpInterfaceMethods<ArithFastMathInterface>],
                   traits)>,
     Arguments<(ins FloatLike:$operand,
       DefaultValuedAttr<Arith_FastMathAttr, "FastMathFlags::none">:$fastmath)>,
     Results<(outs FloatLike:$result)> {
   let assemblyFormat = [{ $operand (`fastmath` `` $fastmath^)?
                           attr-dict `:` type($result) }];
 }
 
 // Base class for floating point binary operations.
 class Arith_FloatBinaryOp<string mnemonic, list<Trait> traits = []> :
     Arith_BinaryOp<mnemonic,
-      !listconcat([DeclareOpInterfaceMethods<ArithFastMathInterface>],
+      !listconcat([Pure, DeclareOpInterfaceMethods<ArithFastMathInterface>],
                   traits)>,
     Arguments<(ins FloatLike:$lhs, FloatLike:$rhs,
       DefaultValuedAttr<Arith_FastMathAttr, "FastMathFlags::none">:$fastmath)>,
     Results<(outs FloatLike:$result)> {
   let assemblyFormat = [{ $lhs `,` $rhs (`fastmath` `` $fastmath^)?
                           attr-dict `:` type($result) }];
 }
 
 // Base class for arithmetic cast operations. Requires a single operand and
 // result. If either is a shaped type, then the other must be of the same shape.
 class Arith_CastOp<string mnemonic, TypeConstraint From, TypeConstraint To,
                    list<Trait> traits = []> :
-    Arith_Op<mnemonic, traits # [SameOperandsAndResultShape,
+    Arith_Op<mnemonic, traits # [Pure, SameOperandsAndResultShape,
       DeclareOpInterfaceMethods<CastOpInterface>]>,
     Arguments<(ins From:$in)>,
     Results<(outs To:$out)> {
   let assemblyFormat = "$in attr-dict `:` type($in) `to` type($out)";
 }
 
 // Casts do not accept indices. Type constraint for signless-integer-like types
 // excluding indices: signless integers, vectors or tensors thereof.

 // Cast from a floating point type to another floating point type.
 class Arith_FToFCastOp<string mnemonic, list<Trait> traits = []> :
     Arith_CastOp<mnemonic, FloatLike, FloatLike, traits>;
 
 // Base class for compare operations. Requires two operands of the same type
 // and returns a single `BoolLike` result. If the operand type is a vector or
 // tensor, then the result will be one of `i1` of the same shape.
 class Arith_CompareOp<string mnemonic, list<Trait> traits = []> :
-    Arith_Op<mnemonic, traits # [SameTypeOperands, TypesMatchWith<
+    Arith_Op<mnemonic, traits # [Pure, SameTypeOperands, TypesMatchWith<
     "result type has i1 element type and same shape as operands",
     "lhs", "result", "::getI1SameShape($_self)">]> {
   let results = (outs BoolLike:$result);
 
   let assemblyFormat = "$predicate `,` $lhs `,` $rhs attr-dict `:` type($lhs)";
 }
 
 // Just like `Arith_CompareOp` but also admits 0-D vectors. Introduced
def Arith_ConstantOp : Op<Arith_Dialect, "constant",
   let assemblyFormat = "attr-dict $value";
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // AddIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_AddIOp : Arith_IntBinaryOp<"addi", [Commutative]> {
+def Arith_AddIOp : Arith_TotalIntBinaryOp<"addi", [Commutative]> {
   let summary = "integer addition operation";
   let description = [{
     The `addi` operation takes two operands and returns one result, each of
     these is required to be the same type. This type may be an integer scalar
     type, a vector whose element type is integer, or a tensor of integers. It
     has no standard attributes.
 
     Example:
let description = [{
     %x = arith.addi %y, %z : tensor<4x?xi8>
     ```
   }];
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 
-def Arith_AddUICarryOp : Arith_Op<"addui_carry", [Commutative,
+def Arith_AddUICarryOp : Arith_Op<"addui_carry", [Pure, Commutative,
     AllTypesMatch<["lhs", "rhs", "sum"]>]> {
   let summary = "unsigned integer addition operation returning sum and carry";
   let description = [{
     The `addui_carry` operation takes two operands and returns two results: the
     sum (same type as both operands), and the carry (boolean-like). The carry
     value `1` indicates unsigned addition overflow, while indicates `0` no
     overflow.
 
let extraClassDeclaration = [{
     ::llvm::Optional<::llvm::SmallVector<int64_t, 4>> getShapeForUnroll();
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // SubIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_SubIOp : Arith_IntBinaryOp<"subi"> {
+def Arith_SubIOp : Arith_TotalIntBinaryOp<"subi"> {
   let summary = "integer subtraction operation";
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // MulIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_MulIOp : Arith_IntBinaryOp<"muli", [Commutative]> {
+def Arith_MulIOp : Arith_TotalIntBinaryOp<"muli", [Commutative]> {
   let summary = "integer multiplication operation";
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // DivUIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_DivUIOp : Arith_IntBinaryOp<"divui"> {
+def Arith_DivUIOp : Arith_IntBinaryOp<"divui", [ConditionallySpeculatable]> {
   let summary = "unsigned integer division operation";
   let description = [{
     Unsigned integer division. Rounds towards zero. Treats the leading bit as
     the most significant, i.e. for `i16` given two's complement representation,
     `6 / -2 = 6 / (2^16 - 2) = 0`.
 
     Note: the semantics of division by zero is TBD; do NOT assume any specific
     behavior.
 
     Example:
 
     ```mlir
     // Scalar unsigned integer division.
     %a = arith.divui %b, %c : i64
 
     // SIMD vector element-wise division.
     %f = arith.divui %g, %h : vector<4xi32>
 
     // Tensor element-wise integer division.
     %x = arith.divui %y, %z : tensor<4x?xi8>
     ```
   }];
+
+  let extraClassDeclaration = [{
+    /// Interface method for ConditionallySpeculatable.
+    Speculation::Speculatability getSpeculatability();
+  }];
+
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // DivSIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_DivSIOp : Arith_IntBinaryOp<"divsi"> {
+def Arith_DivSIOp : Arith_IntBinaryOp<"divsi", [ConditionallySpeculatable]> {
   let summary = "signed integer division operation";
   let description = [{
     Signed integer division. Rounds towards zero. Treats the leading bit as
     sign, i.e. `6 / -2 = -3`.
 
     Note: the semantics of division by zero or signed division overflow (minimum
     value divided by -1) is TBD; do NOT assume any specific behavior.
 
     Example:
 
     ```mlir
     // Scalar signed integer division.
     %a = arith.divsi %b, %c : i64
 
     // SIMD vector element-wise division.
     %f = arith.divsi %g, %h : vector<4xi32>
 
     // Tensor element-wise integer division.
     %x = arith.divsi %y, %z : tensor<4x?xi8>
     ```
   }];
+
+  let extraClassDeclaration = [{
+    /// Interface method for ConditionallySpeculatable.
+    Speculation::Speculatability getSpeculatability();
+  }];
+
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // CeilDivUIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_CeilDivUIOp : Arith_IntBinaryOp<"ceildivui"> {
+def Arith_CeilDivUIOp : Arith_IntBinaryOp<"ceildivui",
+                                          [ConditionallySpeculatable]> {
   let summary = "unsigned ceil integer division operation";
   let description = [{
     Unsigned integer division. Rounds towards positive infinity. Treats the
     leading bit as the most significant, i.e. for `i16` given two's complement
     representation, `6 / -2 = 6 / (2^16 - 2) = 1`.
 
     Note: the semantics of division by zero is TBD; do NOT assume any specific
     behavior.
 
     Example:
 
     ```mlir
     // Scalar unsigned integer division.
     %a = arith.ceildivui %b, %c : i64
     ```
   }];
+
+  let extraClassDeclaration = [{
+    /// Interface method for ConditionallySpeculatable.
+    Speculation::Speculatability getSpeculatability();
+  }];
+
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // CeilDivSIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_CeilDivSIOp : Arith_IntBinaryOp<"ceildivsi"> {
+def Arith_CeilDivSIOp : Arith_IntBinaryOp<"ceildivsi",
+                                          [ConditionallySpeculatable]> {
   let summary = "signed ceil integer division operation";
   let description = [{
     Signed integer division. Rounds towards positive infinity, i.e. `7 / -2 = -3`.
 
     Note: the semantics of division by zero or signed division overflow (minimum
     value divided by -1) is TBD; do NOT assume any specific behavior.
 
     Example:
 
     ```mlir
     // Scalar signed integer division.
     %a = arith.ceildivsi %b, %c : i64
     ```
   }];
+
+  let extraClassDeclaration = [{
+    /// Interface method for ConditionallySpeculatable.
+    Speculation::Speculatability getSpeculatability();
+  }];
+
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // FloorDivSIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_FloorDivSIOp : Arith_IntBinaryOp<"floordivsi"> {
+def Arith_FloorDivSIOp : Arith_TotalIntBinaryOp<"floordivsi"> {
   let summary = "signed floor integer division operation";
   let description = [{
     Signed integer division. Rounds towards negative infinity, i.e. `5 / -2 = -3`.
 
     Note: the semantics of division by zero or signed division overflow (minimum
     value divided by -1) is TBD; do NOT assume any specific behavior.
 
     Example:
 
     ```mlir
     // Scalar signed integer division.
     %a = arith.floordivsi %b, %c : i64
 
     ```
   }];
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // RemUIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_RemUIOp : Arith_IntBinaryOp<"remui"> {
+def Arith_RemUIOp : Arith_TotalIntBinaryOp<"remui"> {
   let summary = "unsigned integer division remainder operation";
   let description = [{
     Unsigned integer division remainder. Treats the leading bit as the most
     significant, i.e. for `i16`, `6 % -2 = 6 % (2^16 - 2) = 6`.
 
     Note: the semantics of division by zero is TBD; do NOT assume any specific
     behavior.
 
def Arith_RemUIOp : Arith_TotalIntBinaryOp<"remui"> {
   }];
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // RemSIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_RemSIOp : Arith_IntBinaryOp<"remsi"> {
+def Arith_RemSIOp : Arith_TotalIntBinaryOp<"remsi"> {
   let summary = "signed integer division remainder operation";
   let description = [{
     Signed integer division remainder. Treats the leading bit as sign, i.e. `6 %
     -2 = 0`.
 
     Note: the semantics of division by zero is TBD; do NOT assume any specific
     behavior.
 
def Arith_RemSIOp : Arith_TotalIntBinaryOp<"remsi"> {
   }];
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // AndIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_AndIOp : Arith_IntBinaryOp<"andi", [Commutative, Idempotent]> {
+def Arith_AndIOp : Arith_TotalIntBinaryOp<"andi", [Commutative, Idempotent]> {
   let summary = "integer binary and";
   let description = [{
     The `andi` operation takes two operands and returns one result, each of
     these is required to be the same type. This type may be an integer scalar
     type, a vector whose element type is integer, or a tensor of integers. It
     has no standard attributes.
 
     Example:
def Arith_AndIOp : Arith_TotalIntBinaryOp<"andi", [Commutative, Idempotent]> {
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // OrIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_OrIOp : Arith_IntBinaryOp<"ori", [Commutative, Idempotent]> {
+def Arith_OrIOp : Arith_TotalIntBinaryOp<"ori", [Commutative, Idempotent]> {
   let summary = "integer binary or";
   let description = [{
     The `ori` operation takes two operands and returns one result, each of these
     is required to be the same type. This type may be an integer scalar type, a
     vector whose element type is integer, or a tensor of integers. It has no
     standard attributes.
 
     Example:
def Arith_OrIOp : Arith_TotalIntBinaryOp<"ori", [Commutative, Idempotent]> {
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // XOrIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_XOrIOp : Arith_IntBinaryOp<"xori", [Commutative]> {
+def Arith_XOrIOp : Arith_TotalIntBinaryOp<"xori", [Commutative]> {
   let summary = "integer binary xor";
   let description = [{
     The `xori` operation takes two operands and returns one result, each of
     these is required to be the same type. This type may be an integer scalar
     type, a vector whose element type is integer, or a tensor of integers. It
     has no standard attributes.
 
     Example:
def Arith_XOrIOp : Arith_TotalIntBinaryOp<"xori", [Commutative]> {
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // ShLIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_ShLIOp : Arith_IntBinaryOp<"shli"> {
+def Arith_ShLIOp : Arith_TotalIntBinaryOp<"shli"> {
   let summary = "integer left-shift";
   let description = [{
     The `shli` operation shifts an integer value to the left by a variable
     amount. The low order bits are filled with zeros.
 
     Example:
 
     ```mlir
     %1 = arith.constant 5 : i8                 // %1 is 0b00000101
     %2 = arith.constant 3 : i8
     %3 = arith.shli %1, %2 : (i8, i8) -> i8    // %3 is 0b00101000
     ```
   }];
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // ShRUIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_ShRUIOp : Arith_IntBinaryOp<"shrui"> {
+def Arith_ShRUIOp : Arith_TotalIntBinaryOp<"shrui"> {
   let summary = "unsigned integer right-shift";
   let description = [{
     The `shrui` operation shifts an integer value to the right by a variable
     amount. The integer is interpreted as unsigned. The high order bits are
     always filled with zeros.
 
     Example:
 
     ```mlir
     %1 = arith.constant 160 : i8               // %1 is 0b10100000
     %2 = arith.constant 3 : i8
     %3 = arith.shrui %1, %2 : (i8, i8) -> i8   // %3 is 0b00010100
     ```
   }];
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // ShRSIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_ShRSIOp : Arith_IntBinaryOp<"shrsi"> {
+def Arith_ShRSIOp : Arith_TotalIntBinaryOp<"shrsi"> {
   let summary = "signed integer right-shift";
   let description = [{
     The `shrsi` operation shifts an integer value to the right by a variable
     amount. The integer is interpreted as signed. The high order bits in the
     output are filled with copies of the most-significant bit of the shifted
     value (which means that the sign of the value is preserved).
 
     Example:
def Arith_MaxFOp : Arith_FloatBinaryOp<"maxf", [Commutative]> {
   }];
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // MaxSIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_MaxSIOp : Arith_IntBinaryOp<"maxsi", [Commutative]> {
+def Arith_MaxSIOp : Arith_TotalIntBinaryOp<"maxsi", [Commutative]> {
   let summary = "signed integer maximum operation";
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // MaxUIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_MaxUIOp : Arith_IntBinaryOp<"maxui", [Commutative]> {
+def Arith_MaxUIOp : Arith_TotalIntBinaryOp<"maxui", [Commutative]> {
   let summary = "unsigned integer maximum operation";
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // MinFOp
 //===----------------------------------------------------------------------===//
 
def Arith_MinFOp : Arith_FloatBinaryOp<"minf", [Commutative]> {
   }];
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // MinSIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_MinSIOp : Arith_IntBinaryOp<"minsi", [Commutative]> {
+def Arith_MinSIOp : Arith_TotalIntBinaryOp<"minsi", [Commutative]> {
   let summary = "signed integer minimum operation";
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // MinUIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_MinUIOp : Arith_IntBinaryOp<"minui", [Commutative]> {
+def Arith_MinUIOp : Arith_TotalIntBinaryOp<"minui", [Commutative]> {
   let summary = "unsigned integer minimum operation";
   let hasFolder = 1;
 }
 
 
 //===----------------------------------------------------------------------===//
 // MulFOp
 //===----------------------------------------------------------------------===//
def Arith_CmpFOp : Arith_CompareOp<"cmpf"> {
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // SelectOp
 //===----------------------------------------------------------------------===//
 
-def SelectOp : Arith_Op<"select", [
+def SelectOp : Arith_Op<"select", [Pure,
     AllTypesMatch<["true_value", "false_value", "result"]>,
     DeclareOpInterfaceMethods<InferIntRangeInterface>,
   ] # ElementwiseMappable.traits> {
   let summary = "select operation";
   let description = [{
     The `arith.select` operation chooses one value based on a binary condition
     supplied as its first operand. If the value of the first operand is `1`,
     the second operand is chosen, otherwise the third operand is chosen.

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

auto result =
           return a;
         }
         return a.udiv(b);
       });
 
   return div0 ? Attribute() : result;
 }
 
+Speculation::Speculatability arith::DivUIOp::getSpeculatability() {
+  // X / 0 => UB
+  return matchPattern(getRhs(), m_NonZero()) ? Speculation::Speculatable
+                                             : Speculation::NotSpeculatable;
+}
+
 //===----------------------------------------------------------------------===//
 // DivSIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::DivSIOp::fold(ArrayRef<Attribute> operands) {
   // divsi (x, 1) -> x.
   if (matchPattern(getRhs(), m_One()))
     return getLhs();
 
   // Don't fold if it would overflow or if it requires a division by zero.
   bool overflowOrDiv0 = false;
   auto result =
       constFoldBinaryOp<IntegerAttr>(operands, [&](APInt a, const APInt &b) {
         if (overflowOrDiv0 || !b) {
           overflowOrDiv0 = true;
           return a;
         }
         return a.sdiv_ov(b, overflowOrDiv0);
       });
 
   return overflowOrDiv0 ? Attribute() : result;
 }
 
+Speculation::Speculatability arith::DivSIOp::getSpeculatability() {
+  bool mayHaveUB = true;
+
+  APInt constRHS;
+  // X / 0 => UB
+  // INT_MIN / -1 => UB
+  if (matchPattern(getRhs(), m_ConstantInt(&constRHS)))
+    mayHaveUB = constRHS.isAllOnes() || constRHS.isZero();
+
+  return mayHaveUB ? Speculation::NotSpeculatable : Speculation::Speculatable;
+}
+
 //===----------------------------------------------------------------------===//
 // Ceil and floor division folding helpers
 //===----------------------------------------------------------------------===//
 
 static APInt signedCeilNonnegInputs(const APInt &a, const APInt &b,
                                     bool &overflow) {
   // Returns (a-1)/b + 1
   APInt one(a.getBitWidth(), 1, true); // Signed value 1.
auto result =
           return quotient;
         APInt one(a.getBitWidth(), 1, true);
         return quotient.uadd_ov(one, overflowOrDiv0);
       });
 
   return overflowOrDiv0 ? Attribute() : result;
 }
 
+Speculation::Speculatability arith::CeilDivUIOp::getSpeculatability() {
+  // X / 0 => UB
+  return matchPattern(getRhs(), m_NonZero()) ? Speculation::Speculatable
+                                             : Speculation::NotSpeculatable;
+}
+
 //===----------------------------------------------------------------------===//
 // CeilDivSIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::CeilDivSIOp::fold(ArrayRef<Attribute> operands) {
   // ceildivsi (x, 1) -> x.
   if (matchPattern(getRhs(), m_One()))
     return getLhs();
auto result =
         APInt posB = zero.ssub_ov(b, overflowOrDiv0);
         APInt div = a.sdiv_ov(posB, overflowOrDiv0);
         return zero.ssub_ov(div, overflowOrDiv0);
       });
 
   return overflowOrDiv0 ? Attribute() : result;
 }
 
+Speculation::Speculatability arith::CeilDivSIOp::getSpeculatability() {
+  bool mayHaveUB = true;
+
+  APInt constRHS;
+  // X / 0 => UB
+  // INT_MIN / -1 => UB
+  if (matchPattern(getRhs(), m_ConstantInt(&constRHS)))
+    mayHaveUB = constRHS.isAllOnes() || constRHS.isZero();
+
+  return mayHaveUB ? Speculation::NotSpeculatable : Speculation::Speculatable;
+}
+
 //===----------------------------------------------------------------------===//
 // FloorDivSIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::FloorDivSIOp::fold(ArrayRef<Attribute> operands) {
   // floordivsi (x, 1) -> x.
   if (matchPattern(getRhs(), m_One()))
     return getLhs();

mlir/test/Transforms/loop-invariant-code-motion.mlir

// CHECK-LABEL: @speculate_memref_dim_known_rank_known_dim_inbounds
   // CHECK: memref.dim
   // CHECK-NEXT: scf.for
   scf.for %i = %lb to %ub step %step {
     %val = memref.dim %t, %c1 : memref<?x?x?x?xf32>
   }
 
   return
 }
+
+// -----
+
+func.func @no_speculate_divui(
+// CHECK-LABEL: @no_speculate_divui(
+    %num: i32, %denom: i32, %lb: index, %ub: index, %step: index) {
+  scf.for %i = %lb to %ub step %step {
+// CHECK: scf.for
+// CHECK: arith.divui
+    %val = arith.divui %num, %denom : i32
+  }
+
+  return
+}
+
+func.func @no_speculate_divsi(
+// CHECK-LABEL: @no_speculate_divsi(
+    %num: i32, %denom: i32, %lb: index, %ub: index, %step: index) {
+  scf.for %i = %lb to %ub step %step {
+// CHECK: scf.for
+// CHECK: arith.divsi
+    %val = arith.divsi %num, %denom : i32
+  }
+
+  return
+}
+
+func.func @no_speculate_ceildivui(
+// CHECK-LABEL: @no_speculate_ceildivui(
+    %num: i32, %denom: i32, %lb: index, %ub: index, %step: index) {
+  scf.for %i = %lb to %ub step %step {
+// CHECK: scf.for
+// CHECK: arith.ceildivui
+    %val = arith.ceildivui %num, %denom : i32
+  }
+
+  return
+}
+
+func.func @no_speculate_ceildivsi(
+// CHECK-LABEL: @no_speculate_ceildivsi(
+    %num: i32, %denom: i32, %lb: index, %ub: index, %step: index) {
+  scf.for %i = %lb to %ub step %step {
+// CHECK: scf.for
+// CHECK: arith.ceildivsi
+    %val = arith.ceildivsi %num, %denom : i32
+  }
+
+  return
+}
+
+func.func @no_speculate_divui_const(%num: i32, %lb: index, %ub: index, %step: index) {
+// CHECK-LABEL: @no_speculate_divui_const(
+  %c0 = arith.constant 0 : i32
+  scf.for %i = %lb to %ub step %step {
+// CHECK: scf.for
+// CHECK: arith.divui
+    %val = arith.divui %num, %c0 : i32
+  }
+
+  return
+}
+
+func.func @speculate_divui_const(
+// CHECK-LABEL: @speculate_divui_const(
+    %num: i32, %lb: index, %ub: index, %step: index) {
+  %c5 = arith.constant 5 : i32
+// CHECK: arith.divui
+// CHECK: scf.for
+  scf.for %i = %lb to %ub step %step {
+    %val = arith.divui %num, %c5 : i32
+  }
+
+  return
+}
+
+func.func @no_speculate_ceildivui_const(%num: i32, %lb: index, %ub: index, %step: index) {
+// CHECK-LABEL: @no_speculate_ceildivui_const(
+  %c0 = arith.constant 0 : i32
+  scf.for %i = %lb to %ub step %step {
+// CHECK: scf.for
+// CHECK: arith.ceildivui
+    %val = arith.ceildivui %num, %c0 : i32
+  }
+
+  return
+}
+
+func.func @speculate_ceildivui_const(
+// CHECK-LABEL: @speculate_ceildivui_const(
+    %num: i32, %lb: index, %ub: index, %step: index) {
+  %c5 = arith.constant 5 : i32
+// CHECK: arith.ceildivui
+// CHECK: scf.for
+  scf.for %i = %lb to %ub step %step {
+    %val = arith.ceildivui %num, %c5 : i32
+  }
+
+  return
+}
+
+func.func @no_speculate_divsi_const0(
+// CHECK-LABEL: @no_speculate_divsi_const0(
+    %num: i32, %denom: i32, %lb: index, %ub: index, %step: index) {
+  %c0 = arith.constant 0 : i32
+  scf.for %i = %lb to %ub step %step {
+// CHECK: scf.for
+// CHECK: arith.divsi
+    %val = arith.divsi %num, %c0 : i32
+  }
+
+  return
+}
+
+func.func @no_speculate_divsi_const_minus1(
+// CHECK-LABEL: @no_speculate_divsi_const_minus1(
+    %num: i32, %denom: i32, %lb: index, %ub: index, %step: index) {
+  %cm1 = arith.constant -1 : i32
+  scf.for %i = %lb to %ub step %step {
+// CHECK: scf.for
+// CHECK: arith.divsi
+    %val = arith.divsi %num, %cm1 : i32
+  }
+
+  return
+}
+
+func.func @speculate_divsi_const(
+// CHECK-LABEL: @speculate_divsi_const(
+    %num: i32, %denom: i32, %lb: index, %ub: index, %step: index) {
+  %c5 = arith.constant 5 : i32
+  scf.for %i = %lb to %ub step %step {
+// CHECK: arith.divsi
+// CHECK: scf.for
+    %val = arith.divsi %num, %c5 : i32
+  }
+
+  return
+}
+
+func.func @no_speculate_ceildivsi_const0(
+// CHECK-LABEL: @no_speculate_ceildivsi_const0(
+    %num: i32, %denom: i32, %lb: index, %ub: index, %step: index) {
+  %c0 = arith.constant 0 : i32
+  scf.for %i = %lb to %ub step %step {
+// CHECK: scf.for
+// CHECK: arith.ceildivsi
+    %val = arith.ceildivsi %num, %c0 : i32
+  }
+
+  return
+}
+
+func.func @no_speculate_ceildivsi_const_minus1(
+// CHECK-LABEL: @no_speculate_ceildivsi_const_minus1(
+    %num: i32, %denom: i32, %lb: index, %ub: index, %step: index) {
+  %cm1 = arith.constant -1 : i32
+  scf.for %i = %lb to %ub step %step {
+// CHECK: scf.for
+// CHECK: arith.ceildivsi
+    %val = arith.ceildivsi %num, %cm1 : i32
+  }
+
+  return
+}
+
+func.func @speculate_ceildivsi_const(
+// CHECK-LABEL: @speculate_ceildivsi_const(
+    %num: i32, %denom: i32, %lb: index, %ub: index, %step: index) {
+  %c5 = arith.constant 5 : i32
+  scf.for %i = %lb to %ub step %step {
+// CHECK: arith.ceildivsi
+// CHECK: scf.for
+    %val = arith.ceildivsi %num, %c5 : i32
+  }
+
+  return
+}