[mlir][Arith] Add ValueBoundsOpInterface impls

[springerm]

Summary

These ops are useful for unit testing. (They do not fold/canonicalize with affine.apply etc.)

Depends On: D145695

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

update API

[sylvestre.ledru]

@springerm seems that https://github.com/llvm/llvm-project/issues/61979 is caused because of this (or the other).
Could you please revert it ? or fix it quickly
thanks

mlir/lib/Dialect/Linalg/Transforms/ValueBoundsOpInterfaceImpl.cpp

 //===- ValueBoundsOpInterfaceImpl.cpp - Impl. of ValueBoundsOpInterface ---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterfaceImpl.h"
 
 #include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using presburger::IntegerPolyhedron;
 
for (int64_t i = 0, e = subViewOp.getMixedSizes().size(); i < e; ++i) {
       }
     }
     llvm_unreachable("could not find non-rank-reduced dim");
   }
 };
 
 } // namespace
 } // namespace memref
+
+namespace arith {
+namespace {
+
+struct AddIOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<AddIOpInterface, AddIOp> {
+  void populateBoundsForIndexValue(Operation *op, Value value,
+                                   ValueBoundsConstraintSet &cstr) const {
+    auto addIOp = cast<AddIOp>(op);
+    assert(value == addIOp.getResult() && "invalid value");
+
+    cstr.addBound(IntegerPolyhedron::BoundType::EQ, value,
+                  cstr.getExpr(addIOp.getLhs()) +
+                      cstr.getExpr(addIOp.getRhs()));
+  }
+};
+
+struct SubIOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<SubIOpInterface, SubIOp> {
+  void populateBoundsForIndexValue(Operation *op, Value value,
+                                   ValueBoundsConstraintSet &cstr) const {
+    auto subIOp = cast<SubIOp>(op);
+    assert(value == subIOp.getResult() && "invalid value");
+
+    cstr.addBound(IntegerPolyhedron::BoundType::EQ, value,
+                  cstr.getExpr(subIOp.getLhs()) -
+                      cstr.getExpr(subIOp.getRhs()));
+  }
+};
+
+struct MulIOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<MulIOpInterface, MulIOp> {
+  void populateBoundsForIndexValue(Operation *op, Value value,
+                                   ValueBoundsConstraintSet &cstr) const {
+    auto mulIOp = cast<MulIOp>(op);
+    assert(value == mulIOp.getResult() && "invalid value");
+
+    cstr.addBound(IntegerPolyhedron::BoundType::EQ, value,
+                  cstr.getExpr(mulIOp.getLhs()) *
+                      cstr.getExpr(mulIOp.getRhs()));
+  }
+};
+
+} // namespace
+} // namespace arith
 } // namespace mlir
 
 void mlir::linalg::registerValueBoundsOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
     tensor::CastOp::attachInterface<tensor::CastOpInterface>(*ctx);
     tensor::DimOp::attachInterface<tensor::DimOpInterface>(*ctx);
     tensor::EmptyOp::attachInterface<tensor::EmptyOpInterface>(*ctx);
registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
     memref::AllocaOp::attachInterface<
         memref::AllocOpInterface<memref::AllocaOp>>(*ctx);
     memref::CastOp::attachInterface<memref::CastOpInterface>(*ctx);
     memref::DimOp::attachInterface<memref::DimOpInterface>(*ctx);
     memref::GetGlobalOp::attachInterface<memref::GetGlobalOpInterface>(*ctx);
     memref::RankOp::attachInterface<memref::RankOpInterface>(*ctx);
     memref::SubViewOp::attachInterface<memref::SubViewOpInterface>(*ctx);
   });
+
+  registry.addExtension(+[](MLIRContext *ctx, arith::ArithDialect *dialect) {
+    arith::AddIOp::attachInterface<arith::AddIOpInterface>(*ctx);
+    arith::SubIOp::attachInterface<arith::SubIOpInterface>(*ctx);
+    arith::MulIOp::attachInterface<arith::MulIOpInterface>(*ctx);
+  });
 }

mlir/test/Dialect/Linalg/value-bounds-op-interface-impl.mlir

 // CHECK-LABEL: func @memref_subview(
 //  CHECK-SAME:     %[[m:.*]]: memref<?xf32>, %[[sz:.*]]: index
 //       CHECK:   return %[[sz]]
 func.func @memref_subview(%m: memref<?xf32>, %sz: index) -> index {
   %0 = memref.subview %m[2][%sz][1] : memref<?xf32> to memref<?xf32, strided<[1], offset: 2>>
   %1 = "test.reify_bound"(%0) {dim = 0} : (memref<?xf32, strided<[1], offset: 2>>) -> (index)
   return %1 : index
 }
+
+// -----
+
+// CHECK: #[[$map]] = affine_map<()[s0] -> (s0 + 5)>
+// CHECK-LABEL: func @arith_addi(
+//  CHECK-SAME:     %[[a:.*]]: index
+//       CHECK:   %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]]]
+//       CHECK:   return %[[apply]]
+func.func @arith_addi(%a: index) -> index {
+  %0 = arith.constant 5 : index
+  %1 = arith.addi %0, %a : index
+  %2 = "test.reify_bound"(%1) : (index) -> (index)
+  return %2 : index
+}
+
+// -----
+
+// CHECK: #[[$map]] = affine_map<()[s0] -> (-s0 + 5)>
+// CHECK-LABEL: func @arith_subi(
+//  CHECK-SAME:     %[[a:.*]]: index
+//       CHECK:   %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]]]
+//       CHECK:   return %[[apply]]
+func.func @arith_subi(%a: index) -> index {
+  %0 = arith.constant 5 : index
+  %1 = arith.subi %0, %a : index
+  %2 = "test.reify_bound"(%1) : (index) -> (index)
+  return %2 : index
+}
+
+// -----
+
+// CHECK: #[[$map]] = affine_map<()[s0] -> (s0 * 5)>
+// CHECK-LABEL: func @arith_muli(
+//  CHECK-SAME:     %[[a:.*]]: index
+//       CHECK:   %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]]]
+//       CHECK:   return %[[apply]]
+func.func @arith_muli(%a: index) -> index {
+  %0 = arith.constant 5 : index
+  %1 = arith.muli %0, %a : index
+  %2 = "test.reify_bound"(%1) : (index) -> (index)
+  return %2 : index
+}