150 lines · cpp
1//===- ReifyValueBounds.cpp --- Reify value bounds with arith 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 <utility>10 11#include "mlir/Dialect/Arith/Transforms/Transforms.h"12 13#include "mlir/Dialect/Arith/IR/Arith.h"14#include "mlir/Dialect/MemRef/IR/MemRef.h"15#include "mlir/Dialect/Tensor/IR/Tensor.h"16#include "mlir/Interfaces/ValueBoundsOpInterface.h"17 18using namespace mlir;19using namespace mlir::arith;20 21/// Build Arith IR for the given affine map and its operands.22static Value buildArithValue(OpBuilder &b, Location loc, AffineMap map,23 ValueRange operands) {24 assert(map.getNumResults() == 1 && "multiple results not supported yet");25 std::function<Value(AffineExpr)> buildExpr = [&](AffineExpr e) -> Value {26 switch (e.getKind()) {27 case AffineExprKind::Constant:28 return ConstantIndexOp::create(b, loc,29 cast<AffineConstantExpr>(e).getValue());30 case AffineExprKind::DimId:31 return operands[cast<AffineDimExpr>(e).getPosition()];32 case AffineExprKind::SymbolId:33 return operands[cast<AffineSymbolExpr>(e).getPosition() +34 map.getNumDims()];35 case AffineExprKind::Add: {36 auto binaryExpr = cast<AffineBinaryOpExpr>(e);37 return AddIOp::create(b, loc, buildExpr(binaryExpr.getLHS()),38 buildExpr(binaryExpr.getRHS()));39 }40 case AffineExprKind::Mul: {41 auto binaryExpr = cast<AffineBinaryOpExpr>(e);42 return MulIOp::create(b, loc, buildExpr(binaryExpr.getLHS()),43 buildExpr(binaryExpr.getRHS()));44 }45 case AffineExprKind::FloorDiv: {46 auto binaryExpr = cast<AffineBinaryOpExpr>(e);47 return DivSIOp::create(b, loc, buildExpr(binaryExpr.getLHS()),48 buildExpr(binaryExpr.getRHS()));49 }50 case AffineExprKind::CeilDiv: {51 auto binaryExpr = cast<AffineBinaryOpExpr>(e);52 return CeilDivSIOp::create(b, loc, buildExpr(binaryExpr.getLHS()),53 buildExpr(binaryExpr.getRHS()));54 }55 case AffineExprKind::Mod: {56 auto binaryExpr = cast<AffineBinaryOpExpr>(e);57 return RemSIOp::create(b, loc, buildExpr(binaryExpr.getLHS()),58 buildExpr(binaryExpr.getRHS()));59 }60 }61 llvm_unreachable("unsupported AffineExpr kind");62 };63 return buildExpr(map.getResult(0));64}65 66FailureOr<OpFoldResult> mlir::arith::reifyValueBound(67 OpBuilder &b, Location loc, presburger::BoundType type,68 const ValueBoundsConstraintSet::Variable &var,69 ValueBoundsConstraintSet::StopConditionFn stopCondition, bool closedUB) {70 // Compute bound.71 AffineMap boundMap;72 ValueDimList mapOperands;73 if (failed(ValueBoundsConstraintSet::computeBound(74 boundMap, mapOperands, type, var, std::move(stopCondition),75 closedUB)))76 return failure();77 78 // Materialize tensor.dim/memref.dim ops.79 SmallVector<Value> operands;80 for (auto valueDim : mapOperands) {81 Value value = valueDim.first;82 std::optional<int64_t> dim = valueDim.second;83 84 if (!dim.has_value()) {85 // This is an index-typed value.86 assert(value.getType().isIndex() && "expected index type");87 operands.push_back(value);88 continue;89 }90 91 assert(cast<ShapedType>(value.getType()).isDynamicDim(*dim) &&92 "expected dynamic dim");93 if (isa<RankedTensorType>(value.getType())) {94 // A tensor dimension is used: generate a tensor.dim.95 operands.push_back(tensor::DimOp::create(b, loc, value, *dim));96 } else if (isa<MemRefType>(value.getType())) {97 // A memref dimension is used: generate a memref.dim.98 operands.push_back(memref::DimOp::create(b, loc, value, *dim));99 } else {100 llvm_unreachable("cannot generate DimOp for unsupported shaped type");101 }102 }103 104 // Check for special cases where no arith ops are needed.105 if (boundMap.isSingleConstant()) {106 // Bound is a constant: return an IntegerAttr.107 return static_cast<OpFoldResult>(108 b.getIndexAttr(boundMap.getSingleConstantResult()));109 }110 // No arith ops are needed if the bound is a single SSA value.111 if (auto expr = dyn_cast<AffineDimExpr>(boundMap.getResult(0)))112 return static_cast<OpFoldResult>(operands[expr.getPosition()]);113 if (auto expr = dyn_cast<AffineSymbolExpr>(boundMap.getResult(0)))114 return static_cast<OpFoldResult>(115 operands[expr.getPosition() + boundMap.getNumDims()]);116 // General case: build Arith ops.117 return static_cast<OpFoldResult>(buildArithValue(b, loc, boundMap, operands));118}119 120FailureOr<OpFoldResult> mlir::arith::reifyShapedValueDimBound(121 OpBuilder &b, Location loc, presburger::BoundType type, Value value,122 int64_t dim, const ValueBoundsConstraintSet::StopConditionFn &stopCondition,123 bool closedUB) {124 auto reifyToOperands = [&](Value v, std::optional<int64_t> d,125 ValueBoundsConstraintSet &cstr) {126 // We are trying to reify a bound for `value` in terms of the owning op's127 // operands. Construct a stop condition that evaluates to "true" for any SSA128 // value expect for `value`. I.e., the bound will be computed in terms of129 // any SSA values expect for `value`. The first such values are operands of130 // the owner of `value`.131 return v != value;132 };133 return reifyValueBound(b, loc, type, {value, dim},134 stopCondition ? stopCondition : reifyToOperands,135 closedUB);136}137 138FailureOr<OpFoldResult> mlir::arith::reifyIndexValueBound(139 OpBuilder &b, Location loc, presburger::BoundType type, Value value,140 const ValueBoundsConstraintSet::StopConditionFn &stopCondition,141 bool closedUB) {142 auto reifyToOperands = [&](Value v, std::optional<int64_t> d,143 ValueBoundsConstraintSet &cstr) {144 return v != value;145 };146 return reifyValueBound(b, loc, type, value,147 stopCondition ? stopCondition : reifyToOperands,148 closedUB);149}150