111 lines · cpp
1//===- ReifyValueBounds.cpp --- Reify value bounds with affine ops ------*-===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/Affine/Transforms/Transforms.h"10 11#include "mlir/Dialect/Affine/IR/AffineOps.h"12#include "mlir/Dialect/MemRef/IR/MemRef.h"13#include "mlir/Dialect/Tensor/IR/Tensor.h"14#include "mlir/Interfaces/ValueBoundsOpInterface.h"15 16using namespace mlir;17using namespace mlir::affine;18 19FailureOr<OpFoldResult> mlir::affine::reifyValueBound(20 OpBuilder &b, Location loc, presburger::BoundType type,21 const ValueBoundsConstraintSet::Variable &var,22 ValueBoundsConstraintSet::StopConditionFn stopCondition, bool closedUB) {23 // Compute bound.24 AffineMap boundMap;25 ValueDimList mapOperands;26 if (failed(ValueBoundsConstraintSet::computeBound(27 boundMap, mapOperands, type, var, stopCondition, closedUB)))28 return failure();29 30 // Reify bound.31 return affine::materializeComputedBound(b, loc, boundMap, mapOperands);32}33 34OpFoldResult affine::materializeComputedBound(35 OpBuilder &b, Location loc, AffineMap boundMap,36 ArrayRef<std::pair<Value, std::optional<int64_t>>> mapOperands) {37 // Materialize tensor.dim/memref.dim ops.38 SmallVector<Value> operands;39 for (auto valueDim : mapOperands) {40 Value value = valueDim.first;41 std::optional<int64_t> dim = valueDim.second;42 43 if (!dim.has_value()) {44 // This is an index-typed value.45 assert(value.getType().isIndex() && "expected index type");46 operands.push_back(value);47 continue;48 }49 50 assert(cast<ShapedType>(value.getType()).isDynamicDim(*dim) &&51 "expected dynamic dim");52 if (isa<RankedTensorType>(value.getType())) {53 // A tensor dimension is used: generate a tensor.dim.54 operands.push_back(tensor::DimOp::create(b, loc, value, *dim));55 } else if (isa<MemRefType>(value.getType())) {56 // A memref dimension is used: generate a memref.dim.57 operands.push_back(memref::DimOp::create(b, loc, value, *dim));58 } else {59 llvm_unreachable("cannot generate DimOp for unsupported shaped type");60 }61 }62 63 // Simplify and return bound.64 affine::canonicalizeMapAndOperands(&boundMap, &operands);65 // Check for special cases where no affine.apply op is needed.66 if (boundMap.isSingleConstant()) {67 // Bound is a constant: return an IntegerAttr.68 return static_cast<OpFoldResult>(69 b.getIndexAttr(boundMap.getSingleConstantResult()));70 }71 // No affine.apply op is needed if the bound is a single SSA value.72 if (auto expr = dyn_cast<AffineDimExpr>(boundMap.getResult(0)))73 return static_cast<OpFoldResult>(operands[expr.getPosition()]);74 if (auto expr = dyn_cast<AffineSymbolExpr>(boundMap.getResult(0)))75 return static_cast<OpFoldResult>(76 operands[expr.getPosition() + boundMap.getNumDims()]);77 // General case: build affine.apply op.78 return static_cast<OpFoldResult>(79 affine::AffineApplyOp::create(b, loc, boundMap, operands).getResult());80}81 82FailureOr<OpFoldResult> mlir::affine::reifyShapedValueDimBound(83 OpBuilder &b, Location loc, presburger::BoundType type, Value value,84 int64_t dim, ValueBoundsConstraintSet::StopConditionFn stopCondition,85 bool closedUB) {86 auto reifyToOperands = [&](Value v, std::optional<int64_t> d,87 ValueBoundsConstraintSet &cstr) {88 // We are trying to reify a bound for `value` in terms of the owning op's89 // operands. Construct a stop condition that evaluates to "true" for any SSA90 // value except for `value`. I.e., the bound will be computed in terms of91 // any SSA values except for `value`. The first such values are operands of92 // the owner of `value`.93 return v != value;94 };95 return reifyValueBound(b, loc, type, {value, dim},96 stopCondition ? stopCondition : reifyToOperands,97 closedUB);98}99 100FailureOr<OpFoldResult> mlir::affine::reifyIndexValueBound(101 OpBuilder &b, Location loc, presburger::BoundType type, Value value,102 ValueBoundsConstraintSet::StopConditionFn stopCondition, bool closedUB) {103 auto reifyToOperands = [&](Value v, std::optional<int64_t> d,104 ValueBoundsConstraintSet &cstr) {105 return v != value;106 };107 return reifyValueBound(b, loc, type, value,108 stopCondition ? stopCondition : reifyToOperands,109 closedUB);110}111