178 lines · cpp
1//===- ValueBoundsOpInterfaceImpl.cpp - Impl. of ValueBoundsOpInterface ---===//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/Arith/IR/ValueBoundsOpInterfaceImpl.h"10 11#include "mlir/Dialect/Arith/IR/Arith.h"12#include "mlir/Interfaces/ValueBoundsOpInterface.h"13 14using namespace mlir;15 16namespace mlir {17namespace arith {18namespace {19 20struct AddIOpInterface21 : public ValueBoundsOpInterface::ExternalModel<AddIOpInterface, AddIOp> {22 void populateBoundsForIndexValue(Operation *op, Value value,23 ValueBoundsConstraintSet &cstr) const {24 auto addIOp = cast<AddIOp>(op);25 assert(value == addIOp.getResult() && "invalid value");26 27 // Note: `getExpr` has a side effect: it may add a new column to the28 // constraint system. The evaluation order of addition operands is29 // unspecified in C++. To make sure that all compilers produce the exact30 // same results (that can be FileCheck'd), it is important that `getExpr`31 // is called first and assigned to temporary variables, and the addition32 // is performed afterwards.33 AffineExpr lhs = cstr.getExpr(addIOp.getLhs());34 AffineExpr rhs = cstr.getExpr(addIOp.getRhs());35 cstr.bound(value) == lhs + rhs;36 }37};38 39struct ConstantOpInterface40 : public ValueBoundsOpInterface::ExternalModel<ConstantOpInterface,41 ConstantOp> {42 void populateBoundsForIndexValue(Operation *op, Value value,43 ValueBoundsConstraintSet &cstr) const {44 auto constantOp = cast<ConstantOp>(op);45 assert(value == constantOp.getResult() && "invalid value");46 47 if (auto attr = llvm::dyn_cast<IntegerAttr>(constantOp.getValue()))48 cstr.bound(value) == attr.getInt();49 }50};51 52struct SubIOpInterface53 : public ValueBoundsOpInterface::ExternalModel<SubIOpInterface, SubIOp> {54 void populateBoundsForIndexValue(Operation *op, Value value,55 ValueBoundsConstraintSet &cstr) const {56 auto subIOp = cast<SubIOp>(op);57 assert(value == subIOp.getResult() && "invalid value");58 59 AffineExpr lhs = cstr.getExpr(subIOp.getLhs());60 AffineExpr rhs = cstr.getExpr(subIOp.getRhs());61 cstr.bound(value) == lhs - rhs;62 }63};64 65struct MulIOpInterface66 : public ValueBoundsOpInterface::ExternalModel<MulIOpInterface, MulIOp> {67 void populateBoundsForIndexValue(Operation *op, Value value,68 ValueBoundsConstraintSet &cstr) const {69 auto mulIOp = cast<MulIOp>(op);70 assert(value == mulIOp.getResult() && "invalid value");71 72 AffineExpr lhs = cstr.getExpr(mulIOp.getLhs());73 AffineExpr rhs = cstr.getExpr(mulIOp.getRhs());74 cstr.bound(value) == lhs *rhs;75 }76};77 78struct FloorDivSIOpInterface79 : public ValueBoundsOpInterface::ExternalModel<FloorDivSIOpInterface,80 FloorDivSIOp> {81 void populateBoundsForIndexValue(Operation *op, Value value,82 ValueBoundsConstraintSet &cstr) const {83 auto divSIOp = cast<FloorDivSIOp>(op);84 assert(value == divSIOp.getResult() && "invalid value");85 86 AffineExpr lhs = cstr.getExpr(divSIOp.getLhs());87 AffineExpr rhs = cstr.getExpr(divSIOp.getRhs());88 cstr.bound(value) == lhs.floorDiv(rhs);89 }90};91 92struct SelectOpInterface93 : public ValueBoundsOpInterface::ExternalModel<SelectOpInterface,94 SelectOp> {95 96 static void populateBounds(SelectOp selectOp, std::optional<int64_t> dim,97 ValueBoundsConstraintSet &cstr) {98 Value value = selectOp.getResult();99 Value condition = selectOp.getCondition();100 Value trueValue = selectOp.getTrueValue();101 Value falseValue = selectOp.getFalseValue();102 103 if (isa<ShapedType>(condition.getType())) {104 // If the condition is a shaped type, the condition is applied105 // element-wise. All three operands must have the same shape.106 cstr.bound(value)[*dim] == cstr.getExpr(trueValue, dim);107 cstr.bound(value)[*dim] == cstr.getExpr(falseValue, dim);108 cstr.bound(value)[*dim] == cstr.getExpr(condition, dim);109 return;110 }111 112 // Populate constraints for the true/false values (and all values on the113 // backward slice, as long as the current stop condition is not satisfied).114 cstr.populateConstraints(trueValue, dim);115 cstr.populateConstraints(falseValue, dim);116 auto boundsBuilder = cstr.bound(value);117 if (dim)118 boundsBuilder[*dim];119 120 // Compare yielded values.121 // If trueValue <= falseValue:122 // * result <= falseValue123 // * result >= trueValue124 if (cstr.populateAndCompare(125 /*lhs=*/{trueValue, dim},126 ValueBoundsConstraintSet::ComparisonOperator::LE,127 /*rhs=*/{falseValue, dim})) {128 if (dim) {129 cstr.bound(value)[*dim] >= cstr.getExpr(trueValue, dim);130 cstr.bound(value)[*dim] <= cstr.getExpr(falseValue, dim);131 } else {132 cstr.bound(value) >= trueValue;133 cstr.bound(value) <= falseValue;134 }135 }136 // If falseValue <= trueValue:137 // * result <= trueValue138 // * result >= falseValue139 if (cstr.populateAndCompare(140 /*lhs=*/{falseValue, dim},141 ValueBoundsConstraintSet::ComparisonOperator::LE,142 /*rhs=*/{trueValue, dim})) {143 if (dim) {144 cstr.bound(value)[*dim] >= cstr.getExpr(falseValue, dim);145 cstr.bound(value)[*dim] <= cstr.getExpr(trueValue, dim);146 } else {147 cstr.bound(value) >= falseValue;148 cstr.bound(value) <= trueValue;149 }150 }151 }152 153 void populateBoundsForIndexValue(Operation *op, Value value,154 ValueBoundsConstraintSet &cstr) const {155 populateBounds(cast<SelectOp>(op), /*dim=*/std::nullopt, cstr);156 }157 158 void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,159 ValueBoundsConstraintSet &cstr) const {160 populateBounds(cast<SelectOp>(op), dim, cstr);161 }162};163} // namespace164} // namespace arith165} // namespace mlir166 167void mlir::arith::registerValueBoundsOpInterfaceExternalModels(168 DialectRegistry ®istry) {169 registry.addExtension(+[](MLIRContext *ctx, arith::ArithDialect *dialect) {170 arith::AddIOp::attachInterface<arith::AddIOpInterface>(*ctx);171 arith::ConstantOp::attachInterface<arith::ConstantOpInterface>(*ctx);172 arith::SubIOp::attachInterface<arith::SubIOpInterface>(*ctx);173 arith::MulIOp::attachInterface<arith::MulIOpInterface>(*ctx);174 arith::FloorDivSIOp::attachInterface<arith::FloorDivSIOpInterface>(*ctx);175 arith::SelectOp::attachInterface<arith::SelectOpInterface>(*ctx);176 });177}178