71 lines · cpp
1//===- AllocationOpInterfaceImpl.cpp - Impl. of AllocationOpInterface -----===//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/MemRef/Transforms/AllocationOpInterfaceImpl.h"10 11#include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.h"12#include "mlir/Dialect/Bufferization/IR/Bufferization.h"13#include "mlir/Dialect/MemRef/IR/MemRef.h"14#include "mlir/IR/Dialect.h"15#include "mlir/IR/Operation.h"16 17using namespace mlir;18 19namespace {20struct DefaultAllocationInterface21 : public bufferization::AllocationOpInterface::ExternalModel<22 DefaultAllocationInterface, memref::AllocOp> {23 static std::optional<Operation *> buildDealloc(OpBuilder &builder,24 Value alloc) {25 return memref::DeallocOp::create(builder, alloc.getLoc(), alloc)26 .getOperation();27 }28 static std::optional<Value> buildClone(OpBuilder &builder, Value alloc) {29 return bufferization::CloneOp::create(builder, alloc.getLoc(), alloc)30 .getResult();31 }32 static ::mlir::HoistingKind getHoistingKind() {33 return HoistingKind::Loop | HoistingKind::Block;34 }35 static ::std::optional<::mlir::Operation *>36 buildPromotedAlloc(OpBuilder &builder, Value alloc) {37 Operation *definingOp = alloc.getDefiningOp();38 return memref::AllocaOp::create(39 builder, definingOp->getLoc(),40 cast<MemRefType>(definingOp->getResultTypes()[0]),41 definingOp->getOperands(), definingOp->getAttrs());42 }43};44 45struct DefaultAutomaticAllocationHoistingInterface46 : public bufferization::AllocationOpInterface::ExternalModel<47 DefaultAutomaticAllocationHoistingInterface, memref::AllocaOp> {48 static ::mlir::HoistingKind getHoistingKind() { return HoistingKind::Loop; }49};50 51struct DefaultReallocationInterface52 : public bufferization::AllocationOpInterface::ExternalModel<53 DefaultAllocationInterface, memref::ReallocOp> {54 static std::optional<Operation *> buildDealloc(OpBuilder &builder,55 Value realloc) {56 return memref::DeallocOp::create(builder, realloc.getLoc(), realloc)57 .getOperation();58 }59};60} // namespace61 62void mlir::memref::registerAllocationOpInterfaceExternalModels(63 DialectRegistry ®istry) {64 registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {65 memref::AllocOp::attachInterface<DefaultAllocationInterface>(*ctx);66 memref::AllocaOp::attachInterface<67 DefaultAutomaticAllocationHoistingInterface>(*ctx);68 memref::ReallocOp::attachInterface<DefaultReallocationInterface>(*ctx);69 });70}71