[mlir][Arith] ValueBoundsOpInterface: Reify with Arith ops

[springerm]

Summary

This revision adds an additional reifyValueBounds helper that reifies the IR with Arith ops instead of Affine ops. This is needed to support value bounds for integer types different from index in a subsequent revision.

Depends On: D147987

[springerm]

rebase

[dcaballe]

mlir/include/mlir/Dialect/Arith/Transforms/Transforms.h

36

I wonder if we could we make the non-index / index type output transparent to the user, i.e., we could have a single interface and then two internal implementations? Or even a single arith-based implementation and then try to fold ops into maps? Not sure if we have any utility that could help with that.

Otherwise, I think at least we should be more specific about the names... something like reifyIndexValueBound and reifyNonIndexValueBound?

[springerm]

mlir/include/mlir/Dialect/Arith/Transforms/Transforms.h

36

we could have a single interface and then two internal implementations
Not sure what this means. We have a single op interface with two methods (ValueBoundsOpInterface::populateBoundsForIndexValue and ValueBoundsOpInterface::populateBoundsForShapeValueDim).

try to fold ops into maps
What kind of maps you mean? Affine maps?

[dcaballe]

mlir/include/mlir/Dialect/Arith/Transforms/Transforms.h

36

(s/interface/API/ in my comment). We have affine::reifyValueBound and arith::reifyValueBound, right? I'm asking if we should unify these two methods into one (if generating affine maps or arith ops depends on the type of operations analized) or, otherwise, use more specific names for the APIs that clearly state the difference between the arith and the affine implementation. We should also elaborate about the differences between the two in the doc. At this point, it's not easy to understand the difference.

try to fold ops into maps
What kind of maps you mean? Affine maps?

This comment was before seeing the follow-up patches. I know see that you are decomposing the resulting affine map when needed.

[springerm]

address comments

[springerm]

mlir/include/mlir/Dialect/Arith/Transforms/Transforms.h

36

Otherwise, I think at least we should be more specific about the names... something like reifyIndexValueBound and reifyNonIndexValueBound?

I split the function in D147987.

We have affine::reifyValueBound and arith::reifyValueBound, right? I'm asking if we should unify these two methods into one (if generating affine maps or arith ops depends on the type of operations analized) or, otherwise, use more specific names for the APIs that clearly state the difference between the arith and the affine implementation.

The kind of operations that are generated depend on which function is called. arith::... generates Arith dialect ops. affine::... generates Affine dialect ops. Users of arith::... do not have to depend on the Affine dialect. (And vice versa.) I think the namespace should be enough to tell the two apart.

Note: The only thing that these functions do is materializing a computed bound in IR. Both call the same ValueBoundsConstraintSet::computeBound helper function.

[dcaballe]

mlir/include/mlir/Dialect/Arith/Transforms/Transforms.h

36

Good point. I hadn't thought about dependences...

No strong opinion but there is no harm in making the name a bit more explicit. The following code from below:

  reified = arith::reifyShapedValueDimBound(
      rewriter, op->getLoc(), *boundType, value, *dim, stopCondition);
} else {
  reified = reifyShapedValueDimBound(
      rewriter, op->getLoc(), *boundType, value, *dim, stopCondition);

would be confusing for someone without too much context on the implementation.

[springerm]

mlir/include/mlir/Dialect/Arith/Transforms/Transforms.h

36

This will read better once the Affine dialect is wrapped in a namespace affine {} like all the other dialects. Then we will have arith::reifyShapedValueDimBound and affine::reifyShapedValueDimBound.

mlir/include/mlir/Dialect/Arith/Transforms/Transforms.h

+//===- Transforms.h - Arith Transforms --------------------------*- C++ -*-===//
+//
+// 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_TRANSFORMS_TRANSFORMS_H
+#define MLIR_DIALECT_ARITH_TRANSFORMS_TRANSFORMS_H
+
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+#include "mlir/Support/LogicalResult.h"
+
+namespace mlir {
+class Location;
+class OpBuilder;
+class OpFoldResult;
+class Value;
+
+namespace presburger {
+enum class BoundType;
+} // namespace presburger
+
+namespace arith {
+
+/// Reify a bound for the given index-typed value or shape dimension size in
+/// terms of the owning op's operands. `dim` must be `nullopt` if and only if
+/// `value` is index-typed.
+///
+/// By default, lower/equal bounds are closed and upper bounds are open. If
+/// `closedUB` is set to "true", upper bounds are also closed.
+FailureOr<OpFoldResult> reifyValueBound(OpBuilder &b, Location loc,
+                                        presburger::BoundType type, Value value,
+                                        std::optional<int64_t> dim,
+                                        bool closedUB = false);
+
+/// Reify a bound for the given index-typed value or shape dimension size in
+/// terms of SSA values for which `stopCondition` is met. `dim` must be
+/// `nullopt` if and only if `value` is index-typed.
+///
+/// By default, lower/equal bounds are closed and upper bounds are open. If
+/// `closedUB` is set to "true", upper bounds are also closed.
+FailureOr<OpFoldResult>
+reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type,
+                Value value, std::optional<int64_t> dim,
+                ValueBoundsConstraintSet::StopConditionFn stopCondition,
+                bool closedUB = false);
+
+} // namespace arith
+} // namespace mlir
+
+#endif // MLIR_DIALECT_ARITH_TRANSFORMS_TRANSFORMS_H

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

 add_mlir_dialect_library(MLIRArithTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
   EmulateWideInt.cpp
   ExpandOps.cpp
   IntRangeOptimizations.cpp
+  ReifyValueBounds.cpp
   UnsignedWhenEquivalent.cpp
 
   ADDITIONAL_HEADER_DIRS
   {$MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arith/Transforms
 
   DEPENDS
   MLIRArithTransformsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAnalysis
   MLIRArithDialect
   MLIRBufferizationDialect
   MLIRBufferizationTransforms
   MLIRFuncDialect
   MLIRFuncTransforms
   MLIRInferIntRangeInterface
   MLIRIR
   MLIRMemRefDialect
   MLIRPass
+  MLIRTensorDialect
   MLIRTransforms
   MLIRTransformUtils
+  MLIRValueBoundsOpInterface
   MLIRVectorDialect
   )

mlir/lib/Dialect/Arith/Transforms/ReifyValueBounds.cpp

+//===- ReifyValueBounds.cpp --- Reify value bounds with arith ops -------*-===//
+//
+// 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/Transforms/Transforms.h"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+
+using namespace mlir;
+using namespace mlir::arith;
+
+/// Build Arith IR for the given affine map and its operands.
+static Value buildArithValue(OpBuilder &b, Location loc, AffineMap map,
+                             ValueRange operands) {
+  assert(map.getNumResults() == 1 && "multiple results not supported yet");
+  std::function<Value(AffineExpr)> buildExpr = [&](AffineExpr e) -> Value {
+    switch (e.getKind()) {
+    case AffineExprKind::Constant:
+      return b.create<ConstantIndexOp>(loc,
+                                       e.cast<AffineConstantExpr>().getValue());
+    case AffineExprKind::DimId:
+      return operands[e.cast<AffineDimExpr>().getPosition()];
+    case AffineExprKind::SymbolId:
+      return operands[e.cast<AffineSymbolExpr>().getPosition() +
+                      map.getNumDims()];
+    case AffineExprKind::Add: {
+      auto binaryExpr = e.cast<AffineBinaryOpExpr>();
+      return b.create<AddIOp>(loc, buildExpr(binaryExpr.getLHS()),
+                              buildExpr(binaryExpr.getRHS()));
+    }
+    case AffineExprKind::Mul: {
+      auto binaryExpr = e.cast<AffineBinaryOpExpr>();
+      return b.create<MulIOp>(loc, buildExpr(binaryExpr.getLHS()),
+                              buildExpr(binaryExpr.getRHS()));
+    }
+    case AffineExprKind::FloorDiv: {
+      auto binaryExpr = e.cast<AffineBinaryOpExpr>();
+      return b.create<DivSIOp>(loc, buildExpr(binaryExpr.getLHS()),
+                               buildExpr(binaryExpr.getRHS()));
+    }
+    case AffineExprKind::CeilDiv: {
+      auto binaryExpr = e.cast<AffineBinaryOpExpr>();
+      return b.create<CeilDivSIOp>(loc, buildExpr(binaryExpr.getLHS()),
+                                   buildExpr(binaryExpr.getRHS()));
+    }
+    case AffineExprKind::Mod: {
+      auto binaryExpr = e.cast<AffineBinaryOpExpr>();
+      return b.create<RemSIOp>(loc, buildExpr(binaryExpr.getLHS()),
+                               buildExpr(binaryExpr.getRHS()));
+    }
+    default:
+      llvm_unreachable("unknown AffineExpr");
+    }
+  };
+  return buildExpr(map.getResult(0));
+}
+
+FailureOr<OpFoldResult> mlir::arith::reifyValueBound(OpBuilder &b, Location loc,
+                                                     presburger::BoundType type,
+                                                     Value value,
+                                                     std::optional<int64_t> dim,
+                                                     bool closedUB) {
+  auto stopCondition = [&](Value v, std::optional<int64_t> d) {
+    // Reify in terms of SSA values that are different from `value`.
+    return v != value;
+  };
+  return reifyValueBound(b, loc, type, value, dim, stopCondition, closedUB);
+}
+
+FailureOr<OpFoldResult> mlir::arith::reifyValueBound(
+    OpBuilder &b, Location loc, presburger::BoundType type, Value value,
+    std::optional<int64_t> dim,
+    function_ref<bool(Value, std::optional<int64_t>)> stopCondition,
+    bool closedUB) {
+  // Compute bound.
+  AffineMap boundMap;
+  ValueDimList mapOperands;
+  if (failed(ValueBoundsConstraintSet::computeBound(
+          boundMap, mapOperands, type, value, dim, stopCondition, closedUB)))
+    return failure();
+
+  // Materialize tensor.dim/memref.dim ops.
+  SmallVector<Value> operands;
+  for (auto valueDim : mapOperands) {
+    Value value = valueDim.first;
+    std::optional<int64_t> dim = valueDim.second;
+
+    if (!dim.has_value()) {
+      // This is an index-typed value.
+      assert(value.getType().isIndex() && "expected index type");
+      operands.push_back(value);
+      continue;
+    }
+
+    assert(value.getType().cast<ShapedType>().isDynamicDim(*dim) &&
+           "expected dynamic dim");
+    if (value.getType().isa<RankedTensorType>()) {
+      // A tensor dimension is used: generate a tensor.dim.
+      operands.push_back(b.create<tensor::DimOp>(loc, value, *dim));
+    } else if (value.getType().isa<MemRefType>()) {
+      // A memref dimension is used: generate a memref.dim.
+      operands.push_back(b.create<memref::DimOp>(loc, value, *dim));
+    } else {
+      llvm_unreachable("cannot generate DimOp for unsupported shaped type");
+    }
+  }
+
+  // Check for special cases where no arith ops are needed.
+  if (boundMap.isSingleConstant()) {
+    // Bound is a constant: return an IntegerAttr.
+    return static_cast<OpFoldResult>(
+        b.getIndexAttr(boundMap.getSingleConstantResult()));
+  }
+  // No arith ops are needed if the bound is a single SSA value.
+  if (auto expr = boundMap.getResult(0).dyn_cast<AffineDimExpr>())
+    return static_cast<OpFoldResult>(operands[expr.getPosition()]);
+  if (auto expr = boundMap.getResult(0).dyn_cast<AffineSymbolExpr>())
+    return static_cast<OpFoldResult>(
+        operands[expr.getPosition() + boundMap.getNumDims()]);
+  // General case: build Arith ops.
+  return static_cast<OpFoldResult>(buildArithValue(b, loc, boundMap, operands));
+}

mlir/test/Dialect/Affine/value-bounds-reification.mlir

 // RUN: mlir-opt %s -test-affine-reify-value-bounds="reify-to-func-args" \
 // RUN:     -verify-diagnostics -split-input-file | FileCheck %s
 
+// RUN: mlir-opt %s -test-affine-reify-value-bounds="reify-to-func-args use-arith-ops" \
+// RUN:     -verify-diagnostics -split-input-file | FileCheck %s --check-prefix=CHECK-ARITH
+
 // CHECK-LABEL: func @reify_through_chain(
 //  CHECK-SAME:     %[[sz0:.*]]: index, %[[sz2:.*]]: index
 //       CHECK:   %[[c10:.*]] = arith.constant 10 : index
 //       CHECK:   return %[[sz0]], %[[c10]], %[[sz2]]
+
+// CHECK-ARITH-LABEL: func @reify_through_chain(
+//  CHECK-ARITH-SAME:     %[[sz0:.*]]: index, %[[sz2:.*]]: index
+//       CHECK-ARITH:   %[[c10:.*]] = arith.constant 10 : index
+//       CHECK-ARITH:   return %[[sz0]], %[[c10]], %[[sz2]]
 func.func @reify_through_chain(%sz0: index, %sz2: index) -> (index, index, index) {
   %c2 = arith.constant 2 : index
   %0 = tensor.empty(%sz0, %sz2) : tensor<?x10x?xf32>
   %1 = tensor.cast %0 : tensor<?x10x?xf32> to tensor<?x?x?xf32>
   %pos = arith.constant 0 : index
   %f = arith.constant 0.0 : f32
   %2 = tensor.insert %f into %1[%pos, %pos, %pos] : tensor<?x?x?xf32>
   %3 = tensor.dim %2, %c2 : tensor<?x?x?xf32>

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
 
+// RUN: mlir-opt %s -test-affine-reify-value-bounds="use-arith-ops" \
+// RUN:     -verify-diagnostics -split-input-file | \
+// RUN: FileCheck %s --check-prefix=CHECK-ARITH
+
 // 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]]
+
+// CHECK-ARITH-LABEL: func @arith_addi(
+//  CHECK-ARITH-SAME:     %[[a:.*]]: index
+//       CHECK-ARITH:   %[[c5:.*]] = arith.constant 5 : index
+//       CHECK-ARITH:   %[[add:.*]] = arith.addi %[[c5]], %[[a]]
+//       CHECK-ARITH:   %[[c5:.*]] = arith.constant 5 : index
+//       CHECK-ARITH:   %[[add:.*]] = arith.addi %[[a]], %[[c5]]
+//       CHECK-ARITH:   return %[[add]]
 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
 }
 
 // -----

mlir/test/lib/Dialect/Affine/CMakeLists.txt

add_mlir_library(MLIRAffineTransformsTestPasses
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
+  MLIRArithTransforms
   MLIRAffineTransforms
   MLIRAffineUtils
   MLIRIR
   MLIRPass
   MLIRSupport
   MLIRMemRefDialect
   MLIRTensorDialect
   MLIRVectorUtils
   )

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

 //===- TestReifyValueBounds.cpp - Test value bounds reification -----------===//
 //
 // 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/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/Transforms/Transforms.h"
+#include "mlir/Dialect/Arith/Transforms/Transforms.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 #include "mlir/Pass/Pass.h"
 
 #define PASS_NAME "test-affine-reify-value-bounds"
struct TestReifyValueBounds
   }
 
   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)};
+
+  Option<bool> useArithOps{*this, "use-arith-ops",
+                           llvm::cl::desc("Reify with arith dialect ops"),
+                           llvm::cl::init(false)};
 };
 
 } // namespace
 
 FailureOr<BoundType> parseBoundType(std::string type) {
   if (type == "EQ")
     return BoundType::EQ;
   if (type == "LB")
     return BoundType::LB;
   if (type == "UB")
     return BoundType::UB;
   return failure();
 }
 
 /// Look for "test.reify_bound" ops in the input and replace their results with
 /// the reified values.
 static LogicalResult testReifyValueBounds(func::FuncOp funcOp,
-                                          bool reifyToFuncArgs) {
+                                          bool reifyToFuncArgs,
+                                          bool useArithOps) {
   IRRewriter rewriter(funcOp.getContext());
   WalkResult result = funcOp.walk([&](Operation *op) {
     // Look for test.reify_bound ops.
     if (op->getName().getStringRef() == "test.reify_bound") {
       if (op->getNumOperands() != 1 || op->getNumResults() != 1 ||
           !op->getResultTypes()[0].isIndex()) {
         op->emitOpError("invalid op");
         return WalkResult::skip();
if (op->getName().getStringRef() == "test.reify_bound") {
       FailureOr<OpFoldResult> reified = failure();
       if (constant) {
         auto reifiedConst = ValueBoundsConstraintSet::computeConstantBound(
             *boundType, value, dim, /*stopCondition=*/nullptr);
         if (succeeded(reifiedConst))
           reified =
               FailureOr<OpFoldResult>(rewriter.getIndexAttr(*reifiedConst));
       } else {
+        if (useArithOps) {
+          reified = arith::reifyValueBound(rewriter, op->getLoc(), *boundType,
+                                           value, dim, stopCondition);
+        } else {
           reified = reifyValueBound(rewriter, op->getLoc(), *boundType, value,
                                     dim, stopCondition);
         }
+      }
       if (failed(reified)) {
         op->emitOpError("could not reify bound");
         return WalkResult::interrupt();
       }
 
       // Replace the op with the reified bound.
       if (auto val = reified->dyn_cast<Value>()) {
         rewriter.replaceOp(op, val);
         return WalkResult::skip();
       }
       Value constOp = rewriter.create<arith::ConstantIndexOp>(
           op->getLoc(), reified->get<Attribute>().cast<IntegerAttr>().getInt());
       rewriter.replaceOp(op, constOp);
       return WalkResult::skip();
     }
     return WalkResult::advance();
   });
   return failure(result.wasInterrupted());
 }
 
 void TestReifyValueBounds::runOnOperation() {
-  if (failed(testReifyValueBounds(getOperation(), reifyToFuncArgs)))
+  if (failed(
+          testReifyValueBounds(getOperation(), reifyToFuncArgs, useArithOps)))
     signalPassFailure();
 }
 
 namespace mlir {
 void registerTestAffineReifyValueBoundsPass() {
   PassRegistration<TestReifyValueBounds>();
 }
 } // namespace mlir

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

deps = [
         ":BufferizationTransforms",
         ":FuncDialect",
         ":FuncTransforms",
         ":IR",
         ":InferIntRangeInterface",
         ":MemRefDialect",
         ":Pass",
         ":Support",
+        ":TensorDialect",
         ":TransformUtils",
         ":Transforms",
+        ":ValueBoundsOpInterface",
         ":VectorDialect",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "ArithUtils",
     srcs = glob([

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

cc_library(
     ]),
     deps = [
         "//llvm:Support",
         "//mlir:AffineAnalysis",
         "//mlir:AffineDialect",
         "//mlir:AffineTransforms",
         "//mlir:AffineUtils",
         "//mlir:Analysis",
+        "//mlir:ArithTransforms",
         "//mlir:DialectUtils",
         "//mlir:FuncDialect",
         "//mlir:IR",
         "//mlir:MemRefDialect",
         "//mlir:Pass",
         "//mlir:SCFDialect",
         "//mlir:Support",
         "//mlir:TensorDialect",