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-exception 6 // 7 //===----------------------------------------------------------------------===// 8 9 #include "mlir/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.h" 10 11 #include "mlir/Dialect/Linalg/IR/Linalg.h" 12 #include "mlir/Interfaces/ValueBoundsOpInterface.h" 13 14 using namespace mlir; 15 16 namespace mlir { 17 namespace linalg { 18 namespace { 19 20 struct IndexOpInterface 21 : public ValueBoundsOpInterface::ExternalModel<IndexOpInterface, IndexOp> { populateBoundsForIndexValuemlir::linalg::__anonda375b9b0111::IndexOpInterface22 void populateBoundsForIndexValue(Operation *op, Value value, 23 ValueBoundsConstraintSet &cstr) const { 24 auto indexOp = cast<IndexOp>(op); 25 auto linalgOp = indexOp->getParentOfType<LinalgOp>(); 26 assert(value == indexOp.getResult() && "invalid value"); 27 28 // index >= 0 29 cstr.bound(value) >= 0; 30 31 // index < dim size 32 int64_t flatDimPos = 33 cast<AffineDimExpr>( 34 linalgOp.getShapesToLoopsMap().getResult(indexOp.getDim())) 35 .getPosition(); 36 // Find the `flatDimPos`-th operand dimension. 37 int64_t flatDimCtr = 0; 38 for (Value operand : linalgOp->getOperands()) { 39 assert(flatDimPos >= flatDimCtr && "invalid pos"); 40 auto shapedType = llvm::cast<ShapedType>(operand.getType()); 41 if (flatDimPos < flatDimCtr + shapedType.getRank()) { 42 cstr.bound(value) < cstr.getExpr(operand, flatDimPos - flatDimCtr); 43 break; 44 } 45 flatDimCtr += shapedType.getRank(); 46 } 47 } 48 }; 49 50 } // namespace 51 } // namespace linalg 52 } // namespace mlir 53 registerValueBoundsOpInterfaceExternalModels(DialectRegistry & registry)54void mlir::linalg::registerValueBoundsOpInterfaceExternalModels( 55 DialectRegistry ®istry) { 56 registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) { 57 IndexOp::attachInterface<IndexOpInterface>(*ctx); 58 // Note: ValueBoundsOpInterface implementation is not required for ops that 59 // implement `DestinationStyleOpInterface` (for querying shaped OpResults). 60 }); 61 } 62