mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.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-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> {
  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 >= 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
  54 void mlir::linalg::registerValueBoundsOpInterfaceExternalModels(
  55     DialectRegistry &registry) {
  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 }