83 lines · cpp
1//===- SubsetOpInterfaceImpl.cpp - Tensor subsets -------------------------===//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/Vector/Transforms/SubsetOpInterfaceImpl.h"10 11#include "mlir/Dialect/Vector/IR/VectorOps.h"12#include "mlir/Interfaces/SubsetOpInterface.h"13 14using namespace mlir;15using namespace mlir::vector;16 17namespace {18 19template <typename OpTy>20struct XferOpSubsetOpInterface21 : public SubsetOpInterface::ExternalModel<XferOpSubsetOpInterface<OpTy>,22 OpTy> {23 FailureOr<HyperrectangularSlice>24 getAccessedHyperrectangularSlice(Operation *op) const {25 auto xferOp = cast<OpTy>(op);26 Builder b(xferOp->getContext());27 SmallVector<OpFoldResult> offsets = llvm::map_to_vector(28 xferOp.getIndices(), [](Value v) -> OpFoldResult { return v; });29 SmallVector<OpFoldResult> sizes = llvm::map_to_vector(30 xferOp.getTransferChunkAccessed(),31 [&](int64_t sz) -> OpFoldResult { return b.getIndexAttr(sz); });32 return HyperrectangularSlice(offsets, sizes);33 }34};35 36struct TransferReadOpSubsetExtractionOpInterface37 : public SubsetExtractionOpInterface::ExternalModel<38 TransferReadOpSubsetExtractionOpInterface, vector::TransferReadOp> {39 OpOperand &getSourceOperand(Operation *op) const {40 return cast<vector::TransferReadOp>(op).getBaseMutable();41 }42};43 44struct TransferWriteOpSubsetInsertionOpInterface45 : public SubsetInsertionOpInterface::ExternalModel<46 TransferWriteOpSubsetInsertionOpInterface, vector::TransferWriteOp> {47 OpOperand &getSourceOperand(Operation *op) const {48 return cast<vector::TransferWriteOp>(op).getValueToStoreMutable();49 }50 51 OpOperand &getDestinationOperand(Operation *op) const {52 return cast<vector::TransferWriteOp>(op).getBaseMutable();53 }54 55 Value buildSubsetExtraction(Operation *op, OpBuilder &builder,56 Location loc) const {57 // TODO: Implement when needed.58 return Value();59 }60 61 SmallVector<Value>62 getValuesNeededToBuildSubsetExtraction(Operation *op) const {63 // TODO: Implement when needed.64 return {};65 }66};67 68} // namespace69 70void mlir::vector::registerSubsetOpInterfaceExternalModels(71 DialectRegistry ®istry) {72 registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {73 TransferReadOp::attachInterface<XferOpSubsetOpInterface<TransferReadOp>>(74 *ctx);75 TransferReadOp::attachInterface<TransferReadOpSubsetExtractionOpInterface>(76 *ctx);77 TransferWriteOp::attachInterface<XferOpSubsetOpInterface<TransferWriteOp>>(78 *ctx);79 TransferWriteOp::attachInterface<TransferWriteOpSubsetInsertionOpInterface>(80 *ctx);81 });82}83