166 lines · cpp
1//===----------------------------------------------------------------------===//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/Affine/IR/AffineOps.h"10#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"11#include "mlir/Dialect/Bufferization/IR/Bufferization.h"12#include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.h"13#include "mlir/Dialect/MemRef/IR/MemRef.h"14#include "mlir/Dialect/Tensor/IR/Tensor.h"15#include "mlir/IR/BuiltinTypes.h"16#include "mlir/Interfaces/FunctionInterfaces.h"17#include "mlir/Transforms/InliningUtils.h"18 19using namespace mlir;20using namespace mlir::bufferization;21 22#include "mlir/Dialect/Bufferization/IR/BufferizationOpsDialect.cpp.inc"23 24//===----------------------------------------------------------------------===//25// Bufferization Dialect Interfaces26//===----------------------------------------------------------------------===//27 28namespace {29struct BufferizationInlinerInterface : public DialectInlinerInterface {30 using DialectInlinerInterface::DialectInlinerInterface;31 32 /// Operations in Bufferization dialect are always legal to inline.33 bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {34 return true;35 }36};37 38template <typename Tensor>39struct BuiltinTensorExternalModel40 : TensorLikeType::ExternalModel<BuiltinTensorExternalModel<Tensor>,41 Tensor> {42 llvm::FailureOr<BufferLikeType> getBufferType(43 mlir::Type tensor, const BufferizationOptions &options,44 llvm::function_ref<mlir::InFlightDiagnostic()> emitError) const {45 auto tensorType = cast<TensorType>(tensor);46 auto memSpace = options.defaultMemorySpaceFn(tensorType);47 if (!memSpace.has_value())48 return emitError() << "could not infer memory space";49 50 return cast<BufferLikeType>(51 getMemRefType(tensorType, options, /*layout=*/{}, *memSpace));52 }53 54 mlir::LogicalResult verifyCompatibleBufferType(55 mlir::Type tensor, BufferLikeType bufferType,56 llvm::function_ref<mlir::InFlightDiagnostic()> emitError) const {57 auto tensorType = cast<ShapedType>(tensor);58 auto memrefType = cast<ShapedType>(bufferType);59 60 if (tensorType.getShape() != memrefType.getShape())61 return emitError() << "shapes do not match";62 63 if (tensorType.getElementType() != memrefType.getElementType())64 return emitError() << "element types do not match";65 66 return mlir::success();67 }68};69 70template <typename MemRef>71struct BuiltinMemRefExternalModel72 : BufferLikeType::ExternalModel<BuiltinMemRefExternalModel<MemRef>,73 MemRef> {};74} // namespace75 76//===----------------------------------------------------------------------===//77// Bufferization Dialect78//===----------------------------------------------------------------------===//79 80void mlir::bufferization::BufferizationDialect::initialize() {81 addOperations<82#define GET_OP_LIST83#include "mlir/Dialect/Bufferization/IR/BufferizationOps.cpp.inc"84 >();85 addInterfaces<BufferizationInlinerInterface>();86 87 // Note: Unlike with other external models, declaring bufferization's88 // "promised interfaces" in builtins for TensorLike and BufferLike type89 // interfaces is not possible (due to builtins being independent of90 // bufferization). Thus, the compromise is to attach these interfaces directly91 // during dialect initialization.92 RankedTensorType::attachInterface<93 BuiltinTensorExternalModel<RankedTensorType>>(*getContext());94 UnrankedTensorType::attachInterface<95 BuiltinTensorExternalModel<UnrankedTensorType>>(*getContext());96 MemRefType::attachInterface<BuiltinMemRefExternalModel<MemRefType>>(97 *getContext());98 UnrankedMemRefType::attachInterface<99 BuiltinMemRefExternalModel<UnrankedMemRefType>>(*getContext());100}101 102LogicalResult BufferizationDialect::verifyRegionArgAttribute(103 Operation *op, unsigned /*regionIndex*/, unsigned argIndex,104 NamedAttribute attr) {105 if (attr.getName() == kWritableAttrName) {106 if (!llvm::isa<BoolAttr>(attr.getValue())) {107 return op->emitError() << "'" << kWritableAttrName108 << "' is expected to be a boolean attribute";109 }110 if (!isa<FunctionOpInterface>(op))111 return op->emitError() << "expected '" << kWritableAttrName112 << "' to be used on function-like operations";113 if (cast<FunctionOpInterface>(op).isExternal())114 return op->emitError() << "'" << kWritableAttrName115 << "' is invalid on external functions";116 return success();117 }118 if (attr.getName() == kBufferAccessAttrName) {119 if (!llvm::isa<StringAttr>(attr.getValue())) {120 return op->emitError() << "'" << kBufferAccessAttrName121 << "' is expected to be a string attribute";122 }123 StringRef str = llvm::cast<StringAttr>(attr.getValue()).getValue();124 if (str != "none" && str != "read" && str != "write" && str != "read-write")125 return op->emitError()126 << "invalid value for '" << kBufferAccessAttrName << "'";127 if (!isa<FunctionOpInterface>(op))128 return op->emitError() << "expected '" << kBufferAccessAttrName129 << "' to be used on function-like operations";130 return success();131 }132 if (attr.getName() == kBufferLayoutAttrName) {133 if (!llvm::isa<MemRefLayoutAttrInterface>(attr.getValue())) {134 return op->emitError() << "'" << kBufferLayoutAttrName135 << "' is expected to be a memref layout attribute";136 }137 if (!isa<FunctionOpInterface>(op))138 return op->emitError() << "expected '" << kBufferLayoutAttrName139 << "' to be used on function-like operations";140 return success();141 }142 return op->emitError() << "attribute '" << kBufferLayoutAttrName143 << "' not supported as a region arg attribute by the "144 "bufferization dialect";145}146 147LogicalResult148BufferizationDialect::verifyOperationAttribute(Operation *op,149 NamedAttribute attr) {150 using bufferization::BufferizableOpInterface;151 152 if (attr.getName() == kManualDeallocation) {153 if (!mlir::hasEffect<MemoryEffects::Allocate>(op) &&154 !mlir::hasEffect<MemoryEffects::Free>(op))155 return op->emitOpError("attribute '")156 << kManualDeallocation157 << "' can be used only on ops that have an allocation and/or free "158 "side effect";159 return success();160 }161 162 return op->emitError()163 << "attribute '" << attr.getName()164 << "' not supported as an op attribute by the bufferization dialect";165}166