223 lines · cpp
1//===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//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/Arith/Transforms/BufferizableOpInterfaceImpl.h"10#include "mlir/Dialect/Arith/IR/Arith.h"11#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"12#include "mlir/Dialect/Bufferization/Transforms/BufferUtils.h"13#include "mlir/Dialect/MemRef/IR/MemRef.h"14#include "mlir/IR/Attributes.h"15#include "mlir/IR/Dialect.h"16#include "mlir/IR/Operation.h"17 18using namespace mlir;19using namespace mlir::bufferization;20 21namespace {22/// Bufferization of arith.constant. Replace with memref.get_global.23struct ConstantOpInterface24 : public BufferizableOpInterface::ExternalModel<ConstantOpInterface,25 arith::ConstantOp> {26 LogicalResult bufferize(Operation *op, RewriterBase &rewriter,27 const BufferizationOptions &options,28 BufferizationState &state) const {29 auto constantOp = cast<arith::ConstantOp>(op);30 auto type = dyn_cast<RankedTensorType>(constantOp.getType());31 32 // Only ranked tensors are supported.33 if (!type)34 return failure();35 36 Attribute memorySpace;37 if (auto memSpace = options.defaultMemorySpaceFn(type))38 memorySpace = *memSpace;39 else40 return constantOp->emitError("could not infer memory space");41 42 // Only constants inside a module are supported.43 auto moduleOp = constantOp->getParentOfType<ModuleOp>();44 if (!moduleOp)45 return failure();46 47 // Create global memory segment and replace tensor with memref pointing to48 // that memory segment.49 FailureOr<memref::GlobalOp> globalOp =50 getGlobalFor(constantOp, state.getSymbolTables(),51 options.bufferAlignment, memorySpace);52 if (failed(globalOp))53 return failure();54 memref::GlobalOp globalMemref = *globalOp;55 replaceOpWithNewBufferizedOp<memref::GetGlobalOp>(56 rewriter, op, globalMemref.getType(), globalMemref.getName());57 58 return success();59 }60 61 bool isWritable(Operation *op, Value value,62 const AnalysisState &state) const {63 // Memory locations returned by memref::GetGlobalOp may not be written to.64 assert(isa<OpResult>(value));65 return false;66 }67};68 69struct IndexCastOpInterface70 : public BufferizableOpInterface::ExternalModel<IndexCastOpInterface,71 arith::IndexCastOp> {72 bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,73 const AnalysisState &state) const {74 return false;75 }76 77 bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,78 const AnalysisState &state) const {79 return false;80 }81 82 AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,83 const AnalysisState &state) const {84 return {{op->getResult(0), BufferRelation::Equivalent}};85 }86 87 LogicalResult bufferize(Operation *op, RewriterBase &rewriter,88 const BufferizationOptions &options,89 BufferizationState &state) const {90 auto castOp = cast<arith::IndexCastOp>(op);91 auto resultTensorType = cast<TensorType>(castOp.getType());92 93 FailureOr<Value> source =94 getBuffer(rewriter, castOp.getIn(), options, state);95 if (failed(source))96 return failure();97 auto sourceType = cast<BaseMemRefType>(source->getType());98 99 // Result type should have same layout and address space as the source type.100 BaseMemRefType resultType;101 if (auto rankedMemRefType = dyn_cast<MemRefType>(sourceType)) {102 resultType = MemRefType::get(103 rankedMemRefType.getShape(), resultTensorType.getElementType(),104 rankedMemRefType.getLayout(), rankedMemRefType.getMemorySpace());105 } else {106 auto unrankedMemrefType = cast<UnrankedMemRefType>(sourceType);107 resultType = UnrankedMemRefType::get(resultTensorType.getElementType(),108 unrankedMemrefType.getMemorySpace());109 }110 111 replaceOpWithNewBufferizedOp<arith::IndexCastOp>(rewriter, op, resultType,112 *source);113 return success();114 }115};116 117/// Bufferization of arith.select. Just replace the operands.118struct SelectOpInterface119 : public BufferizableOpInterface::ExternalModel<SelectOpInterface,120 arith::SelectOp> {121 bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,122 const AnalysisState &state) const {123 return false;124 }125 126 bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,127 const AnalysisState &state) const {128 return false;129 }130 131 AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,132 const AnalysisState &state) const {133 return {{op->getOpResult(0) /*result*/, BufferRelation::Equivalent,134 /*isDefinite=*/false}};135 }136 137 LogicalResult bufferize(Operation *op, RewriterBase &rewriter,138 const BufferizationOptions &options,139 BufferizationState &state) const {140 auto selectOp = cast<arith::SelectOp>(op);141 Location loc = selectOp.getLoc();142 143 // Elementwise conditions are not supported yet. To bufferize such an op,144 // it could be lowered to an elementwise "linalg.generic" with a new145 // "tensor.empty" out tensor, followed by "empty tensor elimination". Such146 // IR will bufferize.147 if (!selectOp.getCondition().getType().isInteger(1))148 return op->emitOpError("only i1 condition values are supported");149 150 // TODO: It would be more efficient to copy the result of the `select` op151 // instead of its OpOperands. In the worst case, 2 copies are inserted at152 // the moment (one for each tensor). When copying the op result, only one153 // copy would be needed.154 FailureOr<Value> maybeTrueBuffer =155 getBuffer(rewriter, selectOp.getTrueValue(), options, state);156 FailureOr<Value> maybeFalseBuffer =157 getBuffer(rewriter, selectOp.getFalseValue(), options, state);158 if (failed(maybeTrueBuffer) || failed(maybeFalseBuffer))159 return failure();160 Value trueBuffer = *maybeTrueBuffer;161 Value falseBuffer = *maybeFalseBuffer;162 163 // The "true" and the "false" operands must have the same type. If the164 // buffers have different types, they differ only in their layout map. Cast165 // both of them to the most dynamic MemRef type.166 if (trueBuffer.getType() != falseBuffer.getType()) {167 auto targetType = bufferization::detail::asMemRefType(168 bufferization::getBufferType(selectOp.getResult(), options, state));169 if (failed(targetType))170 return failure();171 if (trueBuffer.getType() != *targetType)172 trueBuffer =173 memref::CastOp::create(rewriter, loc, *targetType, trueBuffer);174 if (falseBuffer.getType() != *targetType)175 falseBuffer =176 memref::CastOp::create(rewriter, loc, *targetType, falseBuffer);177 }178 179 replaceOpWithNewBufferizedOp<arith::SelectOp>(180 rewriter, op, selectOp.getCondition(), trueBuffer, falseBuffer);181 return success();182 }183 184 FailureOr<BufferLikeType>185 getBufferType(Operation *op, Value value, const BufferizationOptions &options,186 const BufferizationState &state,187 SmallVector<Value> &invocationStack) const {188 auto selectOp = cast<arith::SelectOp>(op);189 assert(value == selectOp.getResult() && "invalid value");190 auto trueType =191 bufferization::detail::asMemRefType(bufferization::getBufferType(192 selectOp.getTrueValue(), options, state, invocationStack));193 auto falseType =194 bufferization::detail::asMemRefType(bufferization::getBufferType(195 selectOp.getFalseValue(), options, state, invocationStack));196 if (failed(trueType) || failed(falseType))197 return failure();198 if (*trueType == *falseType)199 return cast<BufferLikeType>(*trueType);200 if (trueType->getMemorySpace() != falseType->getMemorySpace())201 return op->emitError("inconsistent memory space on true/false operands");202 203 // If the buffers have different types, they differ only in their layout204 // map.205 auto memrefType = llvm::cast<MemRefType>(*trueType);206 return cast<BufferLikeType>(getMemRefTypeWithFullyDynamicLayout(207 RankedTensorType::get(memrefType.getShape(),208 memrefType.getElementType()),209 memrefType.getMemorySpace()));210 }211};212 213} // namespace214 215void mlir::arith::registerBufferizableOpInterfaceExternalModels(216 DialectRegistry ®istry) {217 registry.addExtension(+[](MLIRContext *ctx, ArithDialect *dialect) {218 ConstantOp::attachInterface<ConstantOpInterface>(*ctx);219 IndexCastOp::attachInterface<IndexCastOpInterface>(*ctx);220 SelectOp::attachInterface<SelectOpInterface>(*ctx);221 });222}223