[Mlir-commits] [mlir] [mlir][arith] Add ValueBoundsOpInterface external models for the arith operations DivUI and DivSI (PR #206514)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Hagai Lev Hacohen (HagaiLevHacohen)

<details>
<summary>Changes</summary>

Add ValueBoundsOpInterface external models for the arith operations DivUI and DivSI.

---
Full diff: https://github.com/llvm/llvm-project/pull/206514.diff


3 Files Affected:

- (modified) mlir/lib/Dialect/Arith/IR/ArithDialect.cpp (+4-4) 
- (modified) mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp (+59) 
- (modified) mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir (+59) 


``````````diff
diff --git a/mlir/lib/Dialect/Arith/IR/ArithDialect.cpp b/mlir/lib/Dialect/Arith/IR/ArithDialect.cpp
index 5f645af274b8b..7745868423b2e 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithDialect.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithDialect.cpp
@@ -55,10 +55,10 @@ void arith::ArithDialect::initialize() {
                            SelectOp>();
   declarePromisedInterfaces<bufferization::BufferizableOpInterface, ConstantOp,
                             IndexCastOp, SelectOp>();
-  declarePromisedInterfaces<ValueBoundsOpInterface, AddIOp, ConstantOp, SubIOp,
-                            MulIOp, SelectOp, FloorDivSIOp, CeilDivSIOp,
-                            MinSIOp, MaxSIOp, MinUIOp, MaxUIOp, RemSIOp,
-                            RemUIOp>();
+  declarePromisedInterfaces<ValueBoundsOpInterface, AddIOp, ConstantOp, DivUIOp,
+                            DivSIOp, SubIOp, MulIOp, SelectOp, FloorDivSIOp,
+                            CeilDivSIOp, MinSIOp, MaxSIOp, MinUIOp, MaxUIOp,
+                            RemSIOp, RemUIOp>();
 }
 
 /// Materialize an integer or floating point constant.
diff --git a/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
index b50eb6719b4ae..501a9587e8133 100644
--- a/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -49,6 +49,63 @@ struct ConstantOpInterface
   }
 };
 
+struct DivUIOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<DivUIOpInterface,
+                                                   arith::DivUIOp> {
+  void populateBoundsForIndexValue(Operation *op, Value value,
+                                   ValueBoundsConstraintSet &cstr) const {
+    auto divOp = cast<arith::DivUIOp>(op);
+    assert(value == divOp.getResult() && "invalid value");
+
+    bool lhsNonNegative =
+        ValueBoundsConstraintSet::isProvablyNonNegative(divOp.getLhs(), cstr);
+    bool rhsPositive =
+        ValueBoundsConstraintSet::isProvablyPositive(divOp.getRhs(), cstr);
+    if (!lhsNonNegative || !rhsPositive)
+      return;
+
+    AffineExpr lhs = cstr.getExpr(divOp.getLhs());
+    AffineExpr rhs = cstr.getExpr(divOp.getRhs());
+    cstr.bound(value) >= 0;
+    cstr.bound(value) == lhs.floorDiv(rhs);
+  }
+};
+
+struct DivSIOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<DivSIOpInterface,
+                                                   arith::DivSIOp> {
+  void populateBoundsForIndexValue(Operation *op, Value value,
+                                   ValueBoundsConstraintSet &cstr) const {
+    auto divOp = cast<arith::DivSIOp>(op);
+    assert(value == divOp.getResult() && "invalid value");
+
+    Value lhsValue = divOp.getLhs();
+    Value rhsValue = divOp.getRhs();
+
+    bool lhsNonNegative =
+        ValueBoundsConstraintSet::isProvablyNonNegative(lhsValue, cstr);
+    bool lhsNonPositive =
+        ValueBoundsConstraintSet::isProvablyNonPositive(lhsValue, cstr);
+    bool rhsPositive =
+        ValueBoundsConstraintSet::isProvablyPositive(rhsValue, cstr);
+    bool rhsNegative =
+        ValueBoundsConstraintSet::isProvablyNegative(rhsValue, cstr);
+    if ((!lhsNonNegative && !lhsNonPositive) || (!rhsPositive && !rhsNegative))
+      return;
+
+    AffineExpr lhs = cstr.getExpr(lhsValue);
+    AffineExpr rhs = cstr.getExpr(rhsValue);
+    if ((lhsNonNegative && rhsPositive) || (lhsNonPositive && rhsNegative)) {
+      cstr.bound(value) == lhs.floorDiv(rhs);
+      cstr.bound(value) >= 0;
+    } else if ((lhsNonPositive && rhsPositive) ||
+               (lhsNonNegative && rhsNegative)) {
+      cstr.bound(value) == lhs.ceilDiv(rhs);
+      cstr.bound(value) <= 0;
+    }
+  }
+};
+
 struct SubIOpInterface
     : public ValueBoundsOpInterface::ExternalModel<SubIOpInterface, SubIOp> {
   void populateBoundsForIndexValue(Operation *op, Value value,
@@ -326,6 +383,8 @@ void mlir::arith::registerValueBoundsOpInterfaceExternalModels(
   registry.addExtension(+[](MLIRContext *ctx, arith::ArithDialect *dialect) {
     arith::AddIOp::attachInterface<arith::AddIOpInterface>(*ctx);
     arith::ConstantOp::attachInterface<arith::ConstantOpInterface>(*ctx);
+    arith::DivUIOp::attachInterface<arith::DivUIOpInterface>(*ctx);
+    arith::DivSIOp::attachInterface<arith::DivSIOpInterface>(*ctx);
     arith::SubIOp::attachInterface<arith::SubIOpInterface>(*ctx);
     arith::MulIOp::attachInterface<arith::MulIOpInterface>(*ctx);
     arith::FloorDivSIOp::attachInterface<arith::FloorDivSIOpInterface>(*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 eb67998c52990..bbb902bbcad15 100644
--- a/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
@@ -136,6 +136,65 @@ func.func @arith_ceildivsi_non_pure(%a: index, %b: index) -> index {
 
 // -----
 
+// CHECK-LABEL: func @arith_divui_constant()
+//       CHECK:   %[[c1:.*]] = arith.constant 1 : index
+//       CHECK:   return %[[c1]]
+func.func @arith_divui_constant() -> index {
+  %c7 = arith.constant 7 : index
+  %c5 = arith.constant 5 : index
+  %0 = arith.divui %c7, %c5 : index
+  %1 = "test.reify_bound"(%0) : (index) -> (index)
+  return %1 : index
+}
+
+// -----
+
+
+// CHECK-LABEL: func @arith_divui_positive_lhs()
+//       CHECK:   %[[c1:.*]] = arith.constant 1 : index
+//       CHECK:   return %[[c1]]
+
+func.func @arith_divui_positive_lhs() -> index {
+  %c7 = arith.constant 7 : index
+  %c5 = arith.constant 5 : index
+  %c1 = arith.constant 1 : index
+  %lhs = arith.maxsi %c7, %c1 : index
+  %0 = arith.divui %lhs, %c5 : index
+  %1 = "test.reify_bound"(%0) {type = "LB", constant} : (index) -> (index)
+  return %1 : index
+}
+
+
+// -----
+
+// CHECK-LABEL: func @arith_divsi_negative_positive()
+//       CHECK:   %[[cm1:.*]] = arith.constant -1 : index
+//       CHECK:   return %[[cm1]]
+func.func @arith_divsi_negative_positive() -> index {
+  %cm7 = arith.constant -7 : index
+  %c5 = arith.constant 5 : index
+  %0 = arith.divsi %cm7, %c5 : index
+  %1 = "test.reify_bound"(%0) : (index) -> (index)
+  return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @arith_divsi_negative_lhs()
+//       CHECK:   %[[cm2:.*]] = arith.constant -2 : index
+//       CHECK:   return %[[cm2]]
+func.func @arith_divsi_negative_lhs() -> index {
+  %c2 = arith.constant 2 : index
+  %cm7 = arith.constant -7 : index
+  %cm5 = arith.constant -5 : index
+  %lhs = arith.minsi %cm5, %cm7 : index
+  %0 = arith.divsi %lhs, %c2 : index
+  %1 = "test.reify_bound"(%0) {type = "UB", constant}: (index) -> (index)
+  return %1 : index
+}
+
+// -----
+
 // CHECK-LABEL: func @arith_remsi_positive_positive()
 //       CHECK:   %[[c0:.*]] = arith.constant 0 : index
 //       CHECK:   %[[c5:.*]] = arith.constant 5 : index

``````````

</details>


https://github.com/llvm/llvm-project/pull/206514