From 80b45cc78a0228455cf910c75cac2fe120bb5adb Mon Sep 17 00:00:00 2001 From: Quinn Dawkins 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 { + void populateBoundsForIndexValue(Operation *op, Value value, + ValueBoundsConstraintSet &cstr) const { + auto divSIOp = cast(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 { @@ -157,6 +170,7 @@ void mlir::arith::registerValueBoundsOpInterfaceExternalModels( arith::ConstantOp::attachInterface(*ctx); arith::SubIOp::attachInterface(*ctx); arith::MulIOp::attachInterface(*ctx); + arith::DivSIOp::attachInterface(*ctx); arith::SelectOp::attachInterface(*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 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 { +struct FloorDivSIOpInterface + : public ValueBoundsOpInterface::ExternalModel { void populateBoundsForIndexValue(Operation *op, Value value, ValueBoundsConstraintSet &cstr) const { - auto divSIOp = cast(op); + auto divSIOp = cast(op); assert(value == divSIOp.getResult() && "invalid value"); AffineExpr lhs = cstr.getExpr(divSIOp.getLhs()); @@ -170,7 +171,7 @@ void mlir::arith::registerValueBoundsOpInterfaceExternalModels( arith::ConstantOp::attachInterface(*ctx); arith::SubIOp::attachInterface(*ctx); arith::MulIOp::attachInterface(*ctx); - arith::DivSIOp::attachInterface(*ctx); + arith::FloorDivSIOp::attachInterface(*ctx); arith::SelectOp::attachInterface(*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)