79 lines · cpp
1//===- SubsetInsertionOpInterfaceImpl.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/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h"10 11#include "mlir/Dialect/Linalg/IR/Linalg.h"12#include "mlir/Interfaces/SubsetOpInterface.h"13 14using namespace mlir;15using namespace mlir::linalg;16 17namespace {18struct LinalgCopyOpSubsetOpInterface19 : public SubsetOpInterface::ExternalModel<LinalgCopyOpSubsetOpInterface,20 linalg::CopyOp> {21 bool operatesOnEquivalentSubset(22 Operation *op, SubsetOpInterface candidate,23 function_ref<bool(Value, Value)> equivalenceFn) const {24 // linalg.copy operates on the entire destination tensor.25 if (auto otherCopyOp = dyn_cast<linalg::CopyOp>(candidate.getOperation()))26 return equivalenceFn(cast<linalg::CopyOp>(op).getOutputs()[0],27 otherCopyOp.getOutputs()[0]);28 // In the absence of an analysis, "false" is a conservative way to implement29 // this interface.30 return false;31 }32 33 bool operatesOnDisjointSubset(34 Operation *op, SubsetOpInterface candidate,35 function_ref<bool(Value, Value)> equivalenceFn) const {36 // In the absence of an analysis, "false" is a conservative way to implement37 // this interface.38 return false;39 }40};41 42struct LinalgCopyOpInterface43 : public SubsetInsertionOpInterface::ExternalModel<LinalgCopyOpInterface,44 linalg::CopyOp> {45 OpOperand &getSourceOperand(Operation *op) const {46 auto copyOp = cast<CopyOp>(op);47 return llvm::getSingleElement(copyOp.getInputsMutable());48 }49 50 bool51 isEquivalentSubset(Operation *op, Value candidate,52 function_ref<bool(Value, Value)> equivalenceFn) const {53 auto copyOp = cast<CopyOp>(op);54 return equivalenceFn(candidate,55 llvm::getSingleElement(copyOp.getOutputs()));56 }57 58 Value buildSubsetExtraction(Operation *op, OpBuilder &builder,59 Location loc) const {60 auto copyOp = cast<CopyOp>(op);61 return llvm::getSingleElement(copyOp.getOutputs());62 }63 64 SmallVector<Value>65 getValuesNeededToBuildSubsetExtraction(Operation *op) const {66 auto copyOp = cast<CopyOp>(op);67 return {llvm::getSingleElement(copyOp.getOutputs())};68 }69};70} // namespace71 72void mlir::linalg::registerSubsetOpInterfaceExternalModels(73 DialectRegistry ®istry) {74 registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {75 linalg::CopyOp::attachInterface<LinalgCopyOpSubsetOpInterface>(*ctx);76 linalg::CopyOp::attachInterface<LinalgCopyOpInterface>(*ctx);77 });78}79