62 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/Linalg/IR/ValueBoundsOpInterfaceImpl.h"10 11#include "mlir/Dialect/Linalg/IR/Linalg.h"12#include "mlir/Interfaces/ValueBoundsOpInterface.h"13 14using namespace mlir;15 16namespace mlir {17namespace linalg {18namespace {19 20struct IndexOpInterface21 : public ValueBoundsOpInterface::ExternalModel<IndexOpInterface, IndexOp> {22 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 >= 029 cstr.bound(value) >= 0;30 31 // index < dim size32 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} // namespace51} // namespace linalg52} // namespace mlir53 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 that59 // implement `DestinationStyleOpInterface` (for querying shaped OpResults).60 });61}62