[mlir] Avoid folding `index.remu` and `index.rems` for 0 rhs

[rikhuijzer]

Summary

As discussed in https://github.com/llvm/llvm-project/issues/59714#issuecomment-1369518768, the folder for the remainder operations should be resillient when the rhs is 0.
The file IndexOps.cpp was already checking for multiple divisions by zero, so I tried to stick to the code style from those checks.

Fixes #59714.

As a side note, is it correct that remainder operations are never optimized away? I would expect that the following code

func.func @remu_test() -> index {
  %c3 = index.constant 2
  %c0 = index.constant 1
  %0 = index.remu %c3, %c0
  return %0 : index
}

would be optimized to

func.func @remu_test() -> index {
  return index.constant 0 : index
}

when called with mlir-opt --convert-scf-to-openmp temp.mlir, but maybe I'm misunderstanding something.

[rikhuijzer]

Thanks for your review, Jeff! Could you maybe also land the patch? I don't have permissions.

[Mogball]

Sure. I can do that.

In D151476#4384652, @rikhuijzer wrote:

As a side note, is it correct that remainder operations are never optimized away? I would expect that the following code

They should get folded, but not necessarily by all conversion passes.

mlir/lib/Dialect/Index/IR/IndexOps.cpp

 
 //===----------------------------------------------------------------------===//
 // RemSOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult RemSOp::fold(FoldAdaptor adaptor) {
   return foldBinaryOpChecked(
       adaptor.getOperands(),
-      [](const APInt &lhs, const APInt &rhs) { return lhs.srem(rhs); });
+      [](const APInt &lhs, const APInt &rhs) -> std::optional<APInt> {
+        // Don't fold division by zero.
+        if (rhs.isZero())
+          return std::nullopt;
+        return lhs.srem(rhs);
+      });
 }
 
 //===----------------------------------------------------------------------===//
 // RemUOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult RemUOp::fold(FoldAdaptor adaptor) {
   return foldBinaryOpChecked(
       adaptor.getOperands(),
-      [](const APInt &lhs, const APInt &rhs) { return lhs.urem(rhs); });
+      [](const APInt &lhs, const APInt &rhs) -> std::optional<APInt> {
+        // Don't fold division by zero.
+        if (rhs.isZero())
+          return std::nullopt;
+        return lhs.urem(rhs);
+      });
 }
 
 //===----------------------------------------------------------------------===//
 // MaxSOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult MaxSOp::fold(FoldAdaptor adaptor) {
   return foldBinaryOpChecked(adaptor.getOperands(),

mlir/test/Dialect/Index/index-canonicalize.mlir

func.func @floordivs_nofold() -> index {
 
   // 32-bit result differs from 64-bit.
   // CHECK: index.floordivs
   %0 = index.floordivs %lhs, %c2
 
   return %0 : index
 }
 
+// CHECK-LABEL: @rems_zerodiv_nofold
+func.func @rems_zerodiv_nofold() -> index {
+  %lhs = index.constant 2
+  %rhs = index.constant 0
+  // CHECK: index.rems
+  %0 = index.rems %lhs, %rhs
+  return %0 : index
+}
+
+// CHECK-LABEL: @remu_zerodiv_nofold
+func.func @remu_zerodiv_nofold() -> index {
+  %lhs = index.constant 2
+  %rhs = index.constant 0
+  // CHECK: index.remu
+  %0 = index.remu %lhs, %rhs
+  return %0 : index
+}
+
 // CHECK-LABEL: @rems
 func.func @rems() -> index {
   %lhs = index.constant -5
   %rhs = index.constant 2
   // CHECK: %[[A:.*]] = index.constant -1
   %0 = index.rems %lhs, %rhs
   // CHECK: return %[[A]]
   return %0 : index