[mlir][Interfaces] ValueBoundsOpInterface: Support LB and UB bounds

[springerm]

Summary

This change also adds support for affine.min and affine.max ops.

Depends On: D145696

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

update API

[springerm]

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

283–288

Note: FlatLinearConstraints sometimes returns an affine_map with multiple results when requesting LB/UB bounds. These cases are not yet supported, so we just assert for now. We can add support when needed. (It won't be difficult, we just have to emit an affine.min/affine.max instead of affine.apply.)

[dcaballe]

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

283–288

Thanks for anticipating my questions! This makes sense.

mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h

 /// Return the sink AffineApplyOp on success or failure if `op` does not
 /// decompose into smaller AffineApplyOps.
 /// Note that this can be undone by canonicalization which tries to
 /// maximally compose chains of AffineApplyOps.
 FailureOr<AffineApplyOp> decompose(RewriterBase &rewriter, AffineApplyOp op);
 
 /// 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.
+/// `value` is index-typed. LB and EQ bounds are closed, UB bounds are open.
 FailureOr<OpFoldResult> reifyValueBound(OpBuilder &b, Location loc,
                                         presburger::BoundType type, Value value,
                                         std::optional<int64_t> dim);
 
 /// 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.
+/// `nullopt` if and only if `value` is index-typed. LB and EQ bounds are
+/// closed, UB bounds are open.
 FailureOr<OpFoldResult>
 reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type,
                 Value value, std::optional<int64_t> dim,
                 function_ref<bool(Value)> stopCondition);
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H

mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp

void populateBoundsForIndexValue(Operation *op, Value value,
     SmallVector<AffineExpr> symReplacements = llvm::to_vector(llvm::map_range(
         applyOp.getSymbolOperands(), [&](Value v) { return cstr.getExpr(v); }));
     AffineExpr bound =
         expr.replaceDimsAndSymbols(dimReplacements, symReplacements);
     cstr.addBound(BoundType::EQ, value, bound);
   };
 };
 
+struct AffineMinOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<AffineMinOpInterface,
+                                                   AffineMinOp> {
+  void populateBoundsForIndexValue(Operation *op, Value value,
+                                   ValueBoundsConstraintSet &cstr) const {
+    auto minOp = cast<AffineMinOp>(op);
+    assert(value == minOp.getResult() && "invalid value");
+
+    // Align affine map results with dims/symbols in the constraint set.
+    for (AffineExpr expr : minOp.getAffineMap().getResults()) {
+      SmallVector<AffineExpr> dimReplacements = llvm::to_vector(llvm::map_range(
+          minOp.getDimOperands(), [&](Value v) { return cstr.getExpr(v); }));
+      SmallVector<AffineExpr> symReplacements = llvm::to_vector(llvm::map_range(
+          minOp.getSymbolOperands(), [&](Value v) { return cstr.getExpr(v); }));
+      AffineExpr bound =
+          expr.replaceDimsAndSymbols(dimReplacements, symReplacements);
+      AffineExpr openBound = bound + 1;
+      cstr.addBound(BoundType::UB, value, openBound);
+    }
+  };
+};
+
+struct AffineMaxOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<AffineMaxOpInterface,
+                                                   AffineMaxOp> {
+  void populateBoundsForIndexValue(Operation *op, Value value,
+                                   ValueBoundsConstraintSet &cstr) const {
+    auto maxOp = cast<AffineMaxOp>(op);
+    assert(value == maxOp.getResult() && "invalid value");
+
+    // Align affine map results with dims/symbols in the constraint set.
+    for (AffineExpr expr : maxOp.getAffineMap().getResults()) {
+      SmallVector<AffineExpr> dimReplacements = llvm::to_vector(llvm::map_range(
+          maxOp.getDimOperands(), [&](Value v) { return cstr.getExpr(v); }));
+      SmallVector<AffineExpr> symReplacements = llvm::to_vector(llvm::map_range(
+          maxOp.getSymbolOperands(), [&](Value v) { return cstr.getExpr(v); }));
+      AffineExpr bound =
+          expr.replaceDimsAndSymbols(dimReplacements, symReplacements);
+      cstr.addBound(BoundType::LB, value, bound);
+    }
+  };
+};
+
 } // namespace
 } // namespace mlir
 
 void mlir::affine::registerValueBoundsOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, AffineDialect *dialect) {
     AffineApplyOp::attachInterface<AffineApplyOpInterface>(*ctx);
+    AffineMaxOp::attachInterface<AffineMaxOpInterface>(*ctx);
+    AffineMinOp::attachInterface<AffineMinOpInterface>(*ctx);
   });
 }

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

if (condition(positionToValueDim[nextPos])) {
       ++nextPos;
     }
   }
 }
 
 LogicalResult ValueBoundsConstraintSet::computeBound(
     AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type,
     Value value, std::optional<int64_t> dim, StopConditionFn stopCondition) {
-  // Only EQ bounds are supported at the moment.
-  assert(type == BoundType::EQ && "unsupported bound type");
-
   Builder b(value.getContext());
   mapOperands.clear();
 
   if (stopCondition(value)) {
     // Special case: If the stop condition is satisfied for the input
     // value/dimension, directly return it.
     mapOperands.push_back(std::make_pair(value, dim));
     resultMap = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0,
cstr.projectOut([&](ValueDim p) {
     return !stopCondition(p.first);
   });
 
   // Compute lower and upper bounds for `valueDim`.
   int64_t pos = cstr.getPos(value, dim.value_or(kIndexValue));
   SmallVector<AffineMap> lb(1), ub(1);
   cstr.cstr.getSliceBounds(pos, 1, value.getContext(), &lb, &ub,
                            /*getClosedUB=*/true);
+
   // Note: There are TODOs in the implementation of `getSliceBounds`. In such a
   // case, no lower/upper bound can be computed at the moment.
-  if (lb.empty() || !lb[0] || ub.empty() || !ub[0] ||
-      lb[0].getNumResults() != 1 || ub[0].getNumResults() != 1)
+  // EQ, UB bounds: upper bound is needed.
+  if ((type != BoundType::LB) &&
+      (ub.empty() || !ub[0] || ub[0].getNumResults() == 0))
     return failure();
+  // EQ, LB bounds: lower bound is needed.
+  if ((type != BoundType::UB) &&
+      (lb.empty() || !lb[0] || lb[0].getNumResults() == 0))
+    return failure();
+
+  // TODO: Generate an affine map with multiple results.
+  if (type != BoundType::LB)
+    assert(ub.size() == 1 && ub[0].getNumResults() == 1 &&
+           "multiple bounds not supported");
+  if (type != BoundType::UB)
+    assert(lb.size() == 1 && lb[0].getNumResults() == 1 &&
+           "multiple bounds not supported");
 
-  // Look for same lower and upper bound: EQ bound.
-  if (ub[0] != lb[0])
+  // EQ bound: lower and upper bound must match.
+  if (type == BoundType::EQ && ub[0] != lb[0])
     return failure();
 
+  AffineMap bound;
+  if (type == BoundType::EQ || type == BoundType::LB) {
+    bound = lb[0];
+  } else {
+    // Computed UB is a closed bound. Turn into an open bound.
+    bound = AffineMap::get(ub[0].getNumDims(), ub[0].getNumSymbols(),
+                           ub[0].getResult(0) + 1);
+  }
+
   // Gather all SSA values that are used in the computed bound.
   assert(cstr.cstr.getNumDimAndSymbolVars() == cstr.positionToValueDim.size() &&
          "inconsistent mapping state");
   SmallVector<AffineExpr> replacementDims, replacementSymbols;
   int64_t numDims = 0, numSymbols = 0;
   for (int64_t i = 0; i < cstr.cstr.getNumDimAndSymbolVars(); ++i) {
     // Skip `value`.
     if (i == pos)
       continue;
     // Check if the position `i` is used in the generated bound. If so, it must
     // be included in the generated affine.apply op.
     bool used = false;
     bool isDim = i < cstr.cstr.getNumDimVars();
     if (isDim) {
-      if (lb[0].isFunctionOfDim(i))
+      if (bound.isFunctionOfDim(i))
         used = true;
     } else {
-      if (lb[0].isFunctionOfSymbol(i - cstr.cstr.getNumDimVars()))
+      if (bound.isFunctionOfSymbol(i - cstr.cstr.getNumDimVars()))
         used = true;
     }
 
     if (!used) {
       // Not used: Remove dim/symbol from the result.
       if (isDim) {
         replacementDims.push_back(b.getAffineConstantExpr(0));
       } else {
if (dim == ValueBoundsConstraintSet::kIndexValue) {
       continue;
     }
 
     assert(value.getType().cast<ShapedType>().isDynamicDim(dim) &&
            "expected dynamic dim");
     mapOperands.push_back(std::make_pair(value, dim));
   }
 
-  resultMap = lb[0].replaceDimsAndSymbols(replacementDims, replacementSymbols,
+  resultMap = bound.replaceDimsAndSymbols(replacementDims, replacementSymbols,
                                           numDims, numSymbols);
   return success();
 }

mlir/test/Dialect/Affine/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, s1] -> (s0 + s1)>
 // CHECK-LABEL: func @affine_apply(
 //  CHECK-SAME:     %[[a:.*]]: index, %[[b:.*]]: index
 //       CHECK:   %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]], %[[b]]]
 //       CHECK:   %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]], %[[b]]]
 //       CHECL:   return %[[apply]]
 func.func @affine_apply(%a: index, %b: index) -> index {
   %0 = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%a, %b]
   %1 = "test.reify_bound"(%0) : (index) -> (index)
   return %1 : index
 }
+
+// -----
+
+// CHECK-LABEL: func @affine_max_lb(
+//  CHECK-SAME:     %[[a:.*]]: index
+//       CHECK:   %[[c2:.*]] = arith.constant 2 : index
+//       CHECK:   return %[[c2]]
+func.func @affine_max_lb(%a: index) -> (index) {
+  // Note: There are two LBs: s0 and 2. FlatAffineValueConstraints always
+  // returns the constant one at the moment.
+  %1 = affine.max affine_map<()[s0] -> (s0, 2)>()[%a]
+  %2 = "test.reify_bound"(%1) {type = "LB"}: (index) -> (index)
+  return %2 : index
+}
+
+// -----
+
+func.func @affine_max_ub(%a: index) -> (index) {
+  %1 = affine.max affine_map<()[s0] -> (s0, 2)>()[%a]
+  // expected-error @below{{could not reify bound}}
+  %2 = "test.reify_bound"(%1) {type = "UB"}: (index) -> (index)
+  return %2 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @affine_min_ub(
+//  CHECK-SAME:     %[[a:.*]]: index
+//       CHECK:   %[[c3:.*]] = arith.constant 3 : index
+//       CHECK:   return %[[c3]]
+func.func @affine_min_ub(%a: index) -> (index) {
+  // Note: There are two UBs: s0 + 1 and 3. FlatAffineValueConstraints always
+  // returns the constant one at the moment.
+  %1 = affine.min affine_map<()[s0] -> (s0, 2)>()[%a]
+  %2 = "test.reify_bound"(%1) {type = "UB"}: (index) -> (index)
+  return %2 : index
+}
+
+// -----
+
+func.func @affine_min_lb(%a: index) -> (index) {
+  %1 = affine.min affine_map<()[s0] -> (s0, 2)>()[%a]
+  // expected-error @below{{could not reify bound}}
+  %2 = "test.reify_bound"(%1) {type = "LB"}: (index) -> (index)
+  return %2 : index
+}

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

 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 #include "mlir/Pass/Pass.h"
 
 #define PASS_NAME "test-affine-reify-value-bounds"
 
 using namespace mlir;
+using mlir::presburger::BoundType;
 
 namespace {
 
 /// This pass applies the permutation on the first maximal perfect nest.
 struct TestReifyValueBounds
     : public PassWrapper<TestReifyValueBounds, OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestReifyValueBounds)
 

 private:
   Option<bool> reifyToFuncArgs{
       *this, "reify-to-func-args",
       llvm::cl::desc("Reify in terms of function args"), 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) {
   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();
       }
       Value value = op->getOperand(0);
       if (value.getType().isa<IndexType>() !=
           !op->hasAttrOfType<IntegerAttr>("dim")) {
         // Op should have "dim" attribute if and only if the operand is an
         // index-typed value.
         op->emitOpError("invalid op");
         return WalkResult::skip();
       }
 
+      // Get bound type.
+      std::string boundTypeStr = "EQ";
+      if (auto boundTypeAttr = op->getAttrOfType<StringAttr>("type"))
+        boundTypeStr = boundTypeAttr.str();
+      auto boundType = parseBoundType(boundTypeStr);
+      if (failed(boundType)) {
+        op->emitOpError("invalid op");
+        return WalkResult::interrupt();
+      }
+
+      // Get shape dimension (if any).
       auto dim = value.getType().isIndex()
                      ? std::nullopt
                      : std::make_optional<int64_t>(
                            op->getAttrOfType<IntegerAttr>("dim").getInt());
 
       // Reify value bound.
       rewriter.setInsertionPointAfter(op);
       FailureOr<OpFoldResult> reified;
       if (!reifyToFuncArgs) {
         // Reify in terms of the op's operands.
-        reified = reifyValueBound(rewriter, op->getLoc(),
-                                  presburger::BoundType::EQ, value, dim);
+        reified =
+            reifyValueBound(rewriter, op->getLoc(), *boundType, value, dim);
       } else {
         // Reify in terms of function block arguments.
         auto stopCondition = [](Value v) {
           auto bbArg = v.dyn_cast<BlockArgument>();
           if (!bbArg)
             return false;
           return isa<FunctionOpInterface>(
               bbArg.getParentBlock()->getParentOp());
         };
-        reified =
-            reifyValueBound(rewriter, op->getLoc(), presburger::BoundType::EQ,
-                            value, dim, stopCondition);
+        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>()) {