[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/include/mlir/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.h

+//===- ValueBoundsOpInterfaceImpl.h - 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_ARITH_IR_VALUEBOUNDSOPINTERFACEIMPL_H
+#define MLIR_DIALECT_ARITH_IR_VALUEBOUNDSOPINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace arith {
+void registerValueBoundsOpInterfaceExternalModels(DialectRegistry &registry);
+} // namespace arith
+} // namespace mlir
+
+#endif // MLIR_DIALECT_ARITH_IR_VALUEBOUNDSOPINTERFACEIMPL_H

mlir/include/mlir/InitAllDialects.h

 #define MLIR_INITALLDIALECTS_H_
 
 #include "mlir/Dialect/AMDGPU/AMDGPUDialect.h"
 #include "mlir/Dialect/AMX/AMXDialect.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"
 #include "mlir/Dialect/ArmSVE/ArmSVEDialect.h"
 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h"
 #include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
inline void registerAllDialects(DialectRegistry &registry) {
   linalg::registerTransformDialectExtension(registry);
   memref::registerTransformDialectExtension(registry);
   scf::registerTransformDialectExtension(registry);
   vector::registerTransformDialectExtension(registry);
 
   // Register all external models.
   affine::registerValueBoundsOpInterfaceExternalModels(registry);
   arith::registerBufferizableOpInterfaceExternalModels(registry);
+  arith::registerValueBoundsOpInterfaceExternalModels(registry);
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);
   linalg::registerTilingInterfaceExternalModels(registry);
   memref::registerBufferizableOpInterfaceExternalModels(registry);
   memref::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
   memref::registerValueBoundsOpInterfaceExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);

mlir/lib/Dialect/Arith/IR/CMakeLists.txt

+set(LLVM_OPTIONAL_SOURCES
+  ArithOps.cpp
+  ArithDialect.cpp
+  InferIntRangeInterfaceImpls.cpp
+  ValueBoundsOpInterfaceImpl.cpp
+  )
+
 set(LLVM_TARGET_DEFINITIONS ArithCanonicalization.td)
 mlir_tablegen(ArithCanonicalization.inc -gen-rewriters)
 add_public_tablegen_target(MLIRArithCanonicalizationIncGen)
 
 add_mlir_dialect_library(MLIRArithDialect
   ArithOps.cpp
   ArithDialect.cpp
   InferIntRangeInterfaceImpls.cpp
+  ValueBoundsOpInterfaceImpl.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arith
 
   DEPENDS
   MLIRArithOpsIncGen
   MLIRArithOpsInterfacesIncGen
 
   LINK_LIBS PUBLIC
   MLIRDialect
   MLIRInferIntRangeCommon
   MLIRInferIntRangeInterface
   MLIRInferTypeOpInterface
   MLIRIR
+  MLIRValueBoundsOpInterface
+  )
+
+add_mlir_dialect_library(MLIRArithValueBoundsOpInterfaceImpl
+  ValueBoundsOpInterfaceImpl.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arith
+
+  LINK_LIBS PUBLIC
+  MLIRArithDialect
+  MLIRIR
+  MLIRValueBoundsOpInterface
   )

mlir/lib/Dialect/Arith/IR/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/Arith/IR/ValueBoundsOpInterfaceImpl.h"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+
+using namespace mlir;
+using presburger::BoundType;
+
+namespace mlir {
+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(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(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(BoundType::EQ, value,
+                  cstr.getExpr(mulIOp.getLhs()) *
+                      cstr.getExpr(mulIOp.getRhs()));
+  }
+};
+
+} // namespace
+} // namespace arith
+} // namespace mlir
+
+void mlir::arith::registerValueBoundsOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  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/Arith/value-bounds-op-interface-impl.mlir

+// RUN: mlir-opt %s -test-affine-reify-value-bounds -verify-diagnostics \
+// RUN:     -split-input-file | FileCheck %s
+
+// 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
+}

mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp

struct TestReifyValueBounds
 
   StringRef getArgument() const final { return PASS_NAME; }
   StringRef getDescription() const final {
     return "Tests affine loop permutation utility";
   }
   TestReifyValueBounds() = default;
   TestReifyValueBounds(const TestReifyValueBounds &pass) : PassWrapper(pass){};
 
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AffineDialect>();
+  }
+
   void runOnOperation() override;
 
 private:
   Option<bool> reifyToFuncArgs{
       *this, "reify-to-func-args",
       llvm::cl::desc("Reify in terms of function args"), llvm::cl::init(false)};
 };
 

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":AffinePassIncGen",
         ":AffineToStandard",
         ":AffineTransformOps",
         ":AffineTransforms",
         ":ArithDialect",
         ":ArithToLLVM",
         ":ArithToSPIRV",
         ":ArithTransforms",
+        ":ArithValueBoundsOpInterfaceImpl",
         ":ArmNeonDialect",
         ":ArmSVEDialect",
         ":ArmSVETransforms",
         ":AsyncDialect",
         ":AsyncPassIncGen",
         ":AsyncToLLVM",
         ":AsyncTransforms",
         ":BufferizationDialect",
deps = [
         ":IR",
         ":Support",
         ":ValueBoundsOpInterfaceIncGen",
         "//llvm:Support",
     ],
 )
 
 cc_library(
+    name = "ArithValueBoundsOpInterfaceImpl",
+    srcs = ["lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp"],
+    hdrs = ["include/mlir/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.h"],
+    includes = ["include"],
+    deps = [
+        ":ArithDialect",
+        ":IR",
+        ":ValueBoundsOpInterface",
+    ],
+)
+
+cc_library(
     name = "TilingInterface",
     srcs = ["lib/Interfaces/TilingInterface.cpp"],
     hdrs = ["include/mlir/Interfaces/TilingInterface.h"],
     includes = ["include"],
     deps = [
         ":DialectUtils",
         ":IR",
         ":Support",
deps = [
         ":FuncTdFiles",
     ],
 )
 
 cc_library(
     name = "ArithDialect",
     srcs = glob(
         [
-            "lib/Dialect/Arith/IR/*.cpp",
-            "lib/Dialect/Arith/IR/*.h",
+            "lib/Dialect/Arith/IR/ArithDialect.cpp",
+            "lib/Dialect/Arith/IR/ArithOps.cpp",
+            "lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp",
         ],
     ),
     hdrs = [
         "include/mlir/Dialect/Arith/IR/Arith.h",
         "include/mlir/Transforms/InliningUtils.h",
     ],
     includes = ["include"],
     deps = [