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-exception
6//
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(result: ValueBoundsConstraintSet::computeBound(
27 resultMap&: 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(Elt: value);
47 continue;
48 }
49
50 assert(cast<ShapedType>(value.getType()).isDynamicDim(*dim) &&
51 "expected dynamic dim");
52 if (isa<RankedTensorType>(Val: value.getType())) {
53 // A tensor dimension is used: generate a tensor.dim.
54 operands.push_back(b.create<tensor::DimOp>(loc, value, *dim));
55 } else if (isa<MemRefType>(Val: value.getType())) {
56 // A memref dimension is used: generate a memref.dim.
57 operands.push_back(b.create<memref::DimOp>(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(map: &boundMap, operands: &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>(Val: boundMap.getResult(idx: 0)))
73 return static_cast<OpFoldResult>(operands[expr.getPosition()]);
74 if (auto expr = dyn_cast<AffineSymbolExpr>(Val: boundMap.getResult(idx: 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 b.create<affine::AffineApplyOp>(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's
89 // operands. Construct a stop condition that evaluates to "true" for any SSA
90 // value except for `value`. I.e., the bound will be computed in terms of
91 // any SSA values except for `value`. The first such values are operands of
92 // the owner of `value`.
93 return v != value;
94 };
95 return reifyValueBound(b, loc, type, var: {value, dim},
96 stopCondition: 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, var: value,
108 stopCondition: stopCondition ? stopCondition : reifyToOperands,
109 closedUB);
110}
111

source code of mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp