[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

261–266 ↗

(On Diff #509966)

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

261–266 ↗

(On Diff #509966)

Thanks for anticipating my questions! This makes sense.

mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h

 /// constraint system and mapping of columns to values/shape dimensions.
 ///
 /// Note: This class maintains its own mapping to SSA values; no SSA values are
 /// mapped in the underlying `FlatAffineValueConstraints`. This is because not
 /// only SSA values but also shape dimensions of SSA values must be mapped.
 class ValueBoundsConstraintSet {
 public:
   /// Reify a bound for the given index-typed value or shape dimension size in
-  /// terms of the owning op's operands.
+  /// terms of the owning op's operands. LB and EQ bounds are closed, UB bounds
+  /// are open.
   static FailureOr<OpFoldResult>
   reifyBound(OpBuilder &b, Location loc,
              presburger::IntegerPolyhedron::BoundType type, Value value,
              int64_t dim = kIndexValue);
 
   /// Reify a bound for the given index-typed value or shape dimension size in
-  /// terms of SSA values for which `stopCondition` is met.
+  /// terms of SSA values for which `stopCondition` is met. LB and EQ bounds are
+  /// closed, UB bounds are open.
   static FailureOr<OpFoldResult>
   reifyBound(OpBuilder &b, Location loc,
              presburger::IntegerPolyhedron::BoundType type, Value value,
              int64_t dim, function_ref<bool(Value)> stopCondition);
 
   /// Dimension indentifier to indicate a value is index-typed.
   static const int64_t kIndexValue = -1;
 

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

 
 FailureOr<OpFoldResult> ValueBoundsConstraintSet::reifyBound(
     OpBuilder &b, Location loc, presburger::IntegerPolyhedron::BoundType type,
     Value value, int64_t dim, function_ref<bool(Value)> stopCondition) {
 #ifndef NDEBUG
   assertValidValueDim(value, dim);
 #endif // NDEBUG
 
-  // Only EQ bounds are supported at the moment.
-  assert(type == presburger::IntegerPolyhedron::BoundType::EQ &&
-         "unsupported bound type");
-
   // Process the backward slice of `value` (i.e., reverse use-def chain) until
   // `stopCondition` is met.
   ValueBoundsConstraintSet cstr(std::make_pair(value, dim));
   int64_t pos = cstr.getPos(value, dim);
   cstr.processWorklist(stopCondition);
 
   // Project out all positions (apart from `pos`) that do not match the stop
   // condition.
if (!stopCondition(cstr.positionToValueDim[nextPos].first)) {
       ++nextPos;
     }
   }
 
   // Compute lower and upper bounds for `value`.
   SmallVector<AffineMap> lb(1), ub(1);
   cstr.cstr.getSliceBounds(pos, 1, b.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 != presburger::IntegerPolyhedron::BoundType::LB) &&
+      (ub.empty() || !ub[0] || ub[0].getNumResults() != 1))
+    return failure();
+  // EQ, LB bounds: lower bound is needed.
+  if ((type != presburger::IntegerPolyhedron::BoundType::UB) &&
+      (lb.empty() || !lb[0] || lb[0].getNumResults() != 1))
     return failure();
 
-  // Look for same lower and upper bound: EQ bound.
-  if (ub[0] != lb[0])
+  // EQ bound: lower and upper bound must match.
+  if (type == presburger::IntegerPolyhedron::BoundType::EQ && ub[0] != lb[0])
     return failure();
 
+  AffineMap bound;
+  if (type == presburger::IntegerPolyhedron::BoundType::EQ ||
+      type == presburger::IntegerPolyhedron::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.
   SmallVector<Value> operands;
   assert(cstr.cstr.getNumDimAndSymbolVars() == cstr.positionToValueDim.size() &&
          "inconsistent mapping state");
   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;
     if (i < cstr.cstr.getNumDimVars()) {
-      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: Put an empty Value (will canonicalize away).
       operands.push_back(Value());
       continue;
     }
for (int64_t i = 0; i < cstr.cstr.getNumDimAndSymbolVars(); ++i) {
     } 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");
     }
   }
 
-  mlir::canonicalizeMapAndOperands(&lb[0], &operands);
+  mlir::canonicalizeMapAndOperands(&bound, &operands);
   // Check for special cases where no affine.apply op is needed.
-  if (lb[0].isSingleConstant()) {
+  if (bound.isSingleConstant()) {
     // Bound is a constant: return an IntegerAttr.
     return static_cast<OpFoldResult>(
-        b.getIndexAttr(lb[0].getSingleConstantResult()));
+        b.getIndexAttr(bound.getSingleConstantResult()));
   }
   // No affine.apply op is needed if the bound is a single SSA value.
-  if (auto expr = lb[0].getResult(0).dyn_cast<AffineDimExpr>())
+  if (auto expr = bound.getResult(0).dyn_cast<AffineDimExpr>())
     return static_cast<OpFoldResult>(operands[expr.getPosition()]);
-  if (auto expr = lb[0].getResult(0).dyn_cast<AffineSymbolExpr>())
+  if (auto expr = bound.getResult(0).dyn_cast<AffineSymbolExpr>())
     return static_cast<OpFoldResult>(
         operands[expr.getPosition() + cstr.cstr.getNumDimVars() - 1]);
   // General case: build affine.apply op.
   return static_cast<OpFoldResult>(
-      b.create<AffineApplyOp>(loc, lb[0], operands).getResult());
+      b.create<AffineApplyOp>(loc, bound, operands).getResult());
 }

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

SmallVector<AffineExpr> dimReplacements = llvm::to_vector(llvm::map_range(
         applyOp.getDimOperands(), [&](Value v) { return cstr.getExpr(v); }));
     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(IntegerPolyhedron::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(IntegerPolyhedron::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(IntegerPolyhedron::BoundType::LB, value, bound);
+    }
+  };
+};
 } // namespace
 
 namespace memref {
 namespace {
 
 template <typename OpTy>
 struct AllocOpInterface
     : public ValueBoundsOpInterface::ExternalModel<AllocOpInterface<OpTy>,
registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
     tensor::InsertSliceOp::attachInterface<
         DstValueBoundsOpInterfaceExternalModel<tensor::InsertSliceOp>>(*ctx);
     tensor::PadOp::attachInterface<tensor::PadOpInterface>(*ctx);
     tensor::RankOp::attachInterface<tensor::RankOpInterface>(*ctx);
   });
 
   registry.addExtension(+[](MLIRContext *ctx, AffineDialect *dialect) {
     AffineApplyOp::attachInterface<AffineApplyOpInterface>(*ctx);
+    AffineMaxOp::attachInterface<AffineMaxOpInterface>(*ctx);
+    AffineMinOp::attachInterface<AffineMinOpInterface>(*ctx);
   });
 
   registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
     memref::AllocOp::attachInterface<memref::AllocOpInterface<memref::AllocOp>>(
         *ctx);
     memref::AllocaOp::attachInterface<
         memref::AllocOpInterface<memref::AllocaOp>>(*ctx);
     memref::CastOp::attachInterface<memref::CastOpInterface>(*ctx);

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

 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
+}
+
+// -----
+
 // CHECK-LABEL: func @memref_alloc(
 //  CHECK-SAME:     %[[sz:.*]]: index
 //       CHECK:   %[[c6:.*]] = arith.constant 6 : index
 //       CHECK:   return %[[c6]], %[[sz]]
 func.func @memref_alloc(%sz: index) -> (index, index) {
   %0 = memref.alloc(%sz) : memref<6x?xf32>
   %1 = "test.reify_bound"(%0) {dim = 0} : (memref<6x?xf32>) -> (index)
   %2 = "test.reify_bound"(%0) {dim = 1} : (memref<6x?xf32>) -> (index)

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

 }
 
 static void applyEraseUnnecessaryInputs(func::FuncOp funcOp) {
   RewritePatternSet patterns(funcOp.getContext());
   populateEraseUnnecessaryInputsPatterns(patterns);
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 }
 
+FailureOr<presburger::IntegerPolyhedron::BoundType>
+parseBoundType(std::string type) {
+  if (type == "EQ")
+    return presburger::IntegerPolyhedron::BoundType::EQ;
+  if (type == "LB")
+    return presburger::IntegerPolyhedron::BoundType::LB;
+  if (type == "UB")
+    return presburger::IntegerPolyhedron::BoundType::UB;
+  return failure();
+}
+
 static LogicalResult reifyValueBounds(func::FuncOp funcOp,
                                       bool reifyToFuncArgs) {
   IRRewriter rewriter(funcOp.getContext());
   WalkResult result = funcOp.walk([&](Operation *op) {
     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();
       }
+      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();
+      }
+
       int64_t dim = value.getType().isIndex()
                         ? ValueBoundsConstraintSet::kIndexValue
                         : op->getAttrOfType<IntegerAttr>("dim").getInt();
 
       rewriter.setInsertionPointAfterValue(value);
       FailureOr<OpFoldResult> reified;
       if (!reifyToFuncArgs) {
-        reified = ValueBoundsConstraintSet::reifyBound(
-            rewriter, op->getLoc(),
-            presburger::IntegerPolyhedron::BoundType::EQ, value, dim);
+        reified = ValueBoundsConstraintSet::reifyBound(rewriter, op->getLoc(),
+                                                       *boundType, value, dim);
       } else {
         auto stopCondition = [](Value v) {
           auto bbArg = v.dyn_cast<BlockArgument>();
           if (!bbArg)
             return false;
           return isa<FunctionOpInterface>(
               bbArg.getParentBlock()->getParentOp());
         };
         reified = ValueBoundsConstraintSet::reifyBound(
-            rewriter, op->getLoc(),
-            presburger::IntegerPolyhedron::BoundType::EQ, value, dim,
-            stopCondition);
+            rewriter, op->getLoc(), *boundType, value, dim, stopCondition);
       }
       if (failed(reified)) {
         op->emitOpError("could not reify bound");
         return WalkResult::interrupt();
       }
       if (auto val = reified->dyn_cast<Value>()) {
         rewriter.replaceOp(op, val);
         return WalkResult::skip();