[Mlir-commits] [mlir] [MLIR][Arith] Add ValueBoundsOpInterface for DivSI (PR #137879)

[Quinn Dawkins]

https://github.com/qedawkins updated https://github.com/llvm/llvm-project/pull/137879

>From 80b45cc78a0228455cf910c75cac2fe120bb5adb Mon Sep 17 00:00:00 2001
From: Quinn Dawkins <quinn at nod-labs.com>
Date: Tue, 29 Apr 2025 17:21:01 -0400
Subject: [PATCH 1/2] [MLIR][Arith] Add ValueBoundsOpInterface for DivSI

---
 .../Arith/IR/ValueBoundsOpInterfaceImpl.cpp   | 14 +++++++++++
 .../Arith/value-bounds-op-interface-impl.mlir | 24 +++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
index 6de151594e3e9..4695ac73455f6 100644
--- a/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -75,6 +75,19 @@ struct MulIOpInterface
   }
 };
 
+struct DivSIOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<DivSIOpInterface, DivSIOp> {
+  void populateBoundsForIndexValue(Operation *op, Value value,
+                                   ValueBoundsConstraintSet &cstr) const {
+    auto divSIOp = cast<DivSIOp>(op);
+    assert(value == divSIOp.getResult() && "invalid value");
+
+    AffineExpr lhs = cstr.getExpr(divSIOp.getLhs());
+    AffineExpr rhs = cstr.getExpr(divSIOp.getRhs());
+    cstr.bound(value) == lhs.floorDiv(rhs);
+  }
+};
+
 struct SelectOpInterface
     : public ValueBoundsOpInterface::ExternalModel<SelectOpInterface,
                                                    SelectOp> {
@@ -157,6 +170,7 @@ void mlir::arith::registerValueBoundsOpInterfaceExternalModels(
     arith::ConstantOp::attachInterface<arith::ConstantOpInterface>(*ctx);
     arith::SubIOp::attachInterface<arith::SubIOpInterface>(*ctx);
     arith::MulIOp::attachInterface<arith::MulIOpInterface>(*ctx);
+    arith::DivSIOp::attachInterface<arith::DivSIOpInterface>(*ctx);
     arith::SelectOp::attachInterface<arith::SelectOpInterface>(*ctx);
   });
 }
diff --git a/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
index 8fb3ba1a1ecce..eb2883c211734 100644
--- a/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
@@ -65,6 +65,30 @@ func.func @arith_muli_non_pure(%a: index, %b: index) -> index {
 
 // -----
 
+// CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 floordiv 5)>
+// CHECK-LABEL: func @arith_divsi(
+//  CHECK-SAME:     %[[a:.*]]: index
+//       CHECK:   %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]]]
+//       CHECK:   return %[[apply]]
+func.func @arith_divsi(%a: index) -> index {
+  %0 = arith.constant 5 : index
+  %1 = arith.divsi %a, %0 : index
+  %2 = "test.reify_bound"(%1) : (index) -> (index)
+  return %2 : index
+}
+
+// -----
+
+func.func @arith_divsi_non_pure(%a: index, %b: index) -> index {
+  %0 = arith.divsi %a, %b : index
+  // Semi-affine expressions (such as "symbol * symbol") are not supported.
+  // expected-error @below{{could not reify bound}}
+  %1 = "test.reify_bound"(%0) : (index) -> (index)
+  return %1 : index
+}
+
+// -----
+
 // CHECK-LABEL: func @arith_const()
 //       CHECK:   %[[c5:.*]] = arith.constant 5 : index
 //       CHECK:   %[[c5:.*]] = arith.constant 5 : index

>From 208d674b335d44e926bf687438f1a79863e7dbc3 Mon Sep 17 00:00:00 2001
From: Quinn Dawkins <quinn at nod-labs.com>
Date: Wed, 30 Apr 2025 09:33:59 -0400
Subject: [PATCH 2/2] divsi -> floordivsi

---
 .../Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp    |  9 +++++----
 .../Dialect/Arith/value-bounds-op-interface-impl.mlir  | 10 +++++-----
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
index 4695ac73455f6..bd0c262f932e3 100644
--- a/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -75,11 +75,12 @@ struct MulIOpInterface
   }
 };
 
-struct DivSIOpInterface
-    : public ValueBoundsOpInterface::ExternalModel<DivSIOpInterface, DivSIOp> {
+struct FloorDivSIOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<FloorDivSIOpInterface,
+                                                   FloorDivSIOp> {
   void populateBoundsForIndexValue(Operation *op, Value value,
                                    ValueBoundsConstraintSet &cstr) const {
-    auto divSIOp = cast<DivSIOp>(op);
+    auto divSIOp = cast<FloorDivSIOp>(op);
     assert(value == divSIOp.getResult() && "invalid value");
 
     AffineExpr lhs = cstr.getExpr(divSIOp.getLhs());
@@ -170,7 +171,7 @@ void mlir::arith::registerValueBoundsOpInterfaceExternalModels(
     arith::ConstantOp::attachInterface<arith::ConstantOpInterface>(*ctx);
     arith::SubIOp::attachInterface<arith::SubIOpInterface>(*ctx);
     arith::MulIOp::attachInterface<arith::MulIOpInterface>(*ctx);
-    arith::DivSIOp::attachInterface<arith::DivSIOpInterface>(*ctx);
+    arith::FloorDivSIOp::attachInterface<arith::FloorDivSIOpInterface>(*ctx);
     arith::SelectOp::attachInterface<arith::SelectOpInterface>(*ctx);
   });
 }
diff --git a/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
index eb2883c211734..ec223c1b026bd 100644
--- a/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
@@ -66,21 +66,21 @@ func.func @arith_muli_non_pure(%a: index, %b: index) -> index {
 // -----
 
 // CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 floordiv 5)>
-// CHECK-LABEL: func @arith_divsi(
+// CHECK-LABEL: func @arith_floordivsi(
 //  CHECK-SAME:     %[[a:.*]]: index
 //       CHECK:   %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]]]
 //       CHECK:   return %[[apply]]
-func.func @arith_divsi(%a: index) -> index {
+func.func @arith_floordivsi(%a: index) -> index {
   %0 = arith.constant 5 : index
-  %1 = arith.divsi %a, %0 : index
+  %1 = arith.floordivsi %a, %0 : index
   %2 = "test.reify_bound"(%1) : (index) -> (index)
   return %2 : index
 }
 
 // -----
 
-func.func @arith_divsi_non_pure(%a: index, %b: index) -> index {
-  %0 = arith.divsi %a, %b : index
+func.func @arith_floordivsi_non_pure(%a: index, %b: index) -> index {
+  %0 = arith.floordivsi %a, %b : index
   // Semi-affine expressions (such as "symbol * symbol") are not supported.
   // expected-error @below{{could not reify bound}}
   %1 = "test.reify_bound"(%0) : (index) -> (index)