552 lines
23 KiB
C++
552 lines
23 KiB
C++
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
|
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
|
|
#include "mlir/IR/BuiltinDialect.h"
|
|
#include "mlir/IR/BuiltinTypeInterfaces.h"
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/Interfaces/FunctionInterfaces.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
|
|
#include "llvm/ADT/SmallSet.h"
|
|
#include "llvm/Support/Casting.h"
|
|
#include "llvm/Support/raw_os_ostream.h"
|
|
|
|
#include <cassert>
|
|
#include <filesystem>
|
|
#include <string>
|
|
#include <utility>
|
|
|
|
#include "Conversion/ONNXToSpatial/ONNXToSpatialCommon.hpp"
|
|
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
|
|
#include "src/Accelerators/PIM/Conversion/SpatialToPim/SpatialToPimCommon.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
|
#include "src/Accelerators/PIM/Pass/PimPasses.hpp"
|
|
#include "src/Compiler/CompilerOptions.hpp"
|
|
|
|
using namespace mlir;
|
|
using namespace onnx_mlir;
|
|
using namespace pim;
|
|
|
|
namespace onnx_mlir {
|
|
|
|
namespace {
|
|
|
|
#include "src/Accelerators/PIM/Conversion/SpatialToPim/SpatialToPim.hpp.inc"
|
|
|
|
struct SpatialToPimPass : PassWrapper<SpatialToPimPass, OperationPass<ModuleOp>> {
|
|
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialToPimPass)
|
|
StringRef getArgument() const override { return "convert-spatial-to-pim"; }
|
|
StringRef getDescription() const override { return "Lower Spatial ops to PIM-ready format"; }
|
|
|
|
SpatialToPimPass() = default;
|
|
SpatialToPimPass(const SpatialToPimPass& pass) {}
|
|
|
|
void runOnOperation() final;
|
|
|
|
private:
|
|
SmallVector<Value> outputTensors;
|
|
size_t coreId = 0;
|
|
SmallVector<Operation*> operationsToRemove;
|
|
|
|
void addResultBuffer(func::ReturnOp& returnOp, IRRewriter& rewriter);
|
|
|
|
void allocateAndInitializeCoreLocalVariables(func::FuncOp funcOp, IRRewriter& rewriter);
|
|
|
|
void runOnReceiveOp(spatial::SpatChannelReceiveOp receiveOp, IRRewriter& rewriter);
|
|
void addReceiveOps(Value& channelSourceOp,
|
|
spatial::SpatChannelNewOp& channel,
|
|
Type& channelTensorType,
|
|
bool& useBroadcastOp,
|
|
IRRewriter& rewriter);
|
|
void replaceBlockArgumentWithRecvOp(spatial::SpatWeightedCompute& computeOp,
|
|
unsigned int argIndex,
|
|
spatial::SpatChannelNewOp& channel,
|
|
Type& tensorType,
|
|
bool useBroadcastOp,
|
|
IRRewriter& rewriter);
|
|
|
|
void runOnComputeOp(spatial::SpatWeightedCompute computeOp, IRRewriter& rewriter);
|
|
|
|
void enlargeVMMOutTensorsToCrossbarSize(func::FuncOp funcOp, IRRewriter& rewriter);
|
|
|
|
void replaceReturnOpOperands(func::ReturnOp& returnOp, IRRewriter& rewriter);
|
|
};
|
|
|
|
} // namespace
|
|
|
|
void SpatialToPimPass::runOnOperation() {
|
|
coreId = 1;
|
|
ModuleOp moduleOp = getOperation();
|
|
MLIRContext* ctx = moduleOp.getContext();
|
|
|
|
ConversionTarget target(*ctx);
|
|
target.addLegalDialect<PimDialect, tensor::TensorDialect, arith::ArithDialect, func::FuncDialect, BuiltinDialect>();
|
|
|
|
RewritePatternSet patterns(ctx);
|
|
populateWithGenerated(patterns);
|
|
|
|
if (failed(applyPartialConversion(moduleOp, target, std::move(patterns)))) {
|
|
signalPassFailure();
|
|
return;
|
|
}
|
|
|
|
auto entryFunc = getPimEntryFunc(moduleOp);
|
|
if (failed(entryFunc)) {
|
|
signalPassFailure();
|
|
return;
|
|
}
|
|
func::FuncOp funcOp = *entryFunc;
|
|
|
|
IRRewriter rewriter(&getContext());
|
|
auto returnOp = cast<func::ReturnOp>(funcOp.front().getTerminator());
|
|
|
|
addResultBuffer(returnOp, rewriter);
|
|
allocateAndInitializeCoreLocalVariables(funcOp, rewriter);
|
|
|
|
for (auto receiveOp : funcOp.getOps<spatial::SpatChannelReceiveOp>()) {
|
|
operationsToRemove.push_back(receiveOp);
|
|
runOnReceiveOp(receiveOp, rewriter);
|
|
}
|
|
for (auto computeOp : funcOp.getOps<spatial::SpatWeightedCompute>()) {
|
|
operationsToRemove.push_back(computeOp);
|
|
runOnComputeOp(computeOp, rewriter);
|
|
}
|
|
|
|
enlargeVMMOutTensorsToCrossbarSize(funcOp, rewriter);
|
|
replaceReturnOpOperands(returnOp, rewriter);
|
|
|
|
// Remove all ComputeOps
|
|
for (auto opToRemove : llvm::reverse(operationsToRemove)) {
|
|
if (!opToRemove->use_empty()) {
|
|
opToRemove->dump();
|
|
for (auto user : opToRemove->getUsers())
|
|
user->dump();
|
|
assert(false && "opToRemove should be unused at this point");
|
|
}
|
|
rewriter.eraseOp(opToRemove);
|
|
}
|
|
|
|
// Dump to file for debug
|
|
dumpModule(moduleOp, "pim");
|
|
}
|
|
|
|
void SpatialToPimPass::runOnComputeOp(spatial::SpatWeightedCompute computeOp, IRRewriter& rewriter) {
|
|
Location loc = computeOp->getLoc();
|
|
|
|
auto& block = computeOp.getRegion().front();
|
|
auto yieldOp = cast<spatial::SpatYieldOp>(block.getTerminator());
|
|
|
|
if (computeOp.getNumResults() != yieldOp.getNumOperands())
|
|
llvm_unreachable("ComputeOp must have same number of results as yieldOp operands");
|
|
|
|
for (auto [result, yieldValue] : llvm::zip(computeOp.getResults(), yieldOp.getOperands())) {
|
|
// If this result has no uses, then just skip it
|
|
if (result.use_empty())
|
|
continue;
|
|
|
|
auto yieldType = cast<TensorType>(yieldValue.getType());
|
|
|
|
/*
|
|
* Here we assume that ReturnOp are only reachable by the following patterns:
|
|
*
|
|
* 1)
|
|
* %0 = spat.compute([...])
|
|
* [%0 has one user, which is a ConcatOp]
|
|
* %1 = tensor.concat(%0)
|
|
* [%1 has one user, which is a ReturnOp]
|
|
* return %1
|
|
*
|
|
* 2)
|
|
* %0 = spat.compute([...])
|
|
* [%0 has one user, which is a ReturnOp]
|
|
* return %0
|
|
*
|
|
* If the IR is like 2), then we can store the tensor to the output global memory location
|
|
*/
|
|
auto resultUses = result.getUses();
|
|
auto numResultUses = rangeLength(resultUses);
|
|
if (numResultUses == 1) {
|
|
OpOperand& resultUse = *resultUses.begin();
|
|
Operation* resultUser = resultUse.getOwner();
|
|
|
|
if (isa<func::ReturnOp>(resultUser)) {
|
|
size_t resultIndexInReturn = resultUse.getOperandNumber();
|
|
size_t offset = 0;
|
|
size_t numElements = yieldType.getNumElements();
|
|
size_t elementSize = yieldType.getElementType().getIntOrFloatBitWidth() / 8;
|
|
|
|
// Store to global memory
|
|
Value outputTensor = outputTensors[resultIndexInReturn];
|
|
rewriter.setInsertionPointAfterValue(yieldValue);
|
|
PimMemCopyDevToHostOp::create(rewriter,
|
|
loc,
|
|
outputTensor.getType(),
|
|
outputTensor,
|
|
yieldValue,
|
|
rewriter.getI32IntegerAttr(offset),
|
|
rewriter.getI32IntegerAttr(0),
|
|
rewriter.getI32IntegerAttr(numElements * elementSize));
|
|
continue;
|
|
}
|
|
|
|
if (isa<tensor::ConcatOp>(resultUser) || isa<spatial::SpatImgConcatOp>(resultUser)) {
|
|
auto concatOp = resultUser;
|
|
auto concatValue = concatOp->getResult(0);
|
|
auto concatUses = concatValue.getUses();
|
|
auto numConcatUses = rangeLength(concatUses);
|
|
if (numConcatUses == 1) {
|
|
OpOperand& concatUse = *concatUses.begin();
|
|
Operation* concatUser = concatUse.getOwner();
|
|
if (isa<func::ReturnOp>(concatUser)) {
|
|
size_t concatIndexInReturn = concatUse.getOperandNumber();
|
|
size_t resultIndexInConcat = resultUses.begin()->getOperandNumber();
|
|
size_t offset = 0;
|
|
for (auto operand : concatOp->getOperands().take_front(resultIndexInConcat))
|
|
offset += cast<ShapedType>(operand.getType()).getNumElements()
|
|
* cast<ShapedType>(operand.getType()).getElementTypeBitWidth() / 8;
|
|
|
|
size_t elementSize = yieldType.getElementTypeBitWidth() / 8;
|
|
|
|
// Store to global memory
|
|
Value outputTensor = outputTensors[concatIndexInReturn];
|
|
rewriter.setInsertionPointAfterValue(yieldValue);
|
|
PimMemCopyDevToHostOp::create(rewriter,
|
|
loc,
|
|
outputTensor.getType(),
|
|
outputTensor,
|
|
yieldValue,
|
|
rewriter.getI32IntegerAttr(offset),
|
|
rewriter.getI32IntegerAttr(0),
|
|
rewriter.getI32IntegerAttr(yieldType.getNumElements() * elementSize));
|
|
continue;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// If this pattern was not found, then create a channel and send the value
|
|
|
|
// 1. Create a new ChannelOp
|
|
rewriter.setInsertionPoint(computeOp);
|
|
auto channelType = spatial::SpatChannelType::get(computeOp.getContext());
|
|
auto channelOp = spatial::SpatChannelNewOp::create(rewriter, loc, channelType);
|
|
|
|
// 2. Receive value through the channel
|
|
// If this result is used by more than one user, then use a "Broadcast"
|
|
// channel operation. However, there is a special case: we have a single
|
|
// user (a ReshapeOp) which in turn is used by multiple ComputeOps. In this
|
|
// case, we need to use a "Broadcast" channel operation. `addReceiveOps`
|
|
// will detect this case and update `useBroadcastOp` accordingly.
|
|
bool useBroadcastOp = (numResultUses > 1);
|
|
addReceiveOps(result, channelOp, yieldType, useBroadcastOp, rewriter);
|
|
|
|
// 3. Send the value through the channel
|
|
rewriter.setInsertionPointAfterValue(yieldValue);
|
|
if (useBroadcastOp)
|
|
spatial::SpatChannelBroadcastSendOp::create(rewriter, loc, channelOp, yieldValue);
|
|
else
|
|
spatial::SpatChannelSendOp::create(rewriter, loc, channelOp, yieldValue);
|
|
}
|
|
|
|
// Use `HaltOp` instead of `YieldOp`
|
|
rewriter.setInsertionPoint(yieldOp);
|
|
rewriter.replaceOpWithNewOp<PimHaltOp>(yieldOp);
|
|
|
|
// Replace `spat.compute` with `pim.core`
|
|
rewriter.setInsertionPointAfter(computeOp);
|
|
auto coreOp = PimCoreOp::create(rewriter, loc, computeOp.getWeights(), rewriter.getI32IntegerAttr(coreId++));
|
|
auto& coreOpBlocks = coreOp.getBody().getBlocks();
|
|
block.eraseArguments(0, block.getNumArguments());
|
|
coreOpBlocks.splice(coreOpBlocks.begin(), computeOp.getBody().getBlocks());
|
|
Block* tempComputeBlock = new Block();
|
|
computeOp.getBody().push_back(tempComputeBlock);
|
|
rewriter.setInsertionPointToEnd(tempComputeBlock);
|
|
PimHaltOp::create(rewriter, computeOp.getLoc());
|
|
}
|
|
|
|
void SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func::FuncOp funcOp, IRRewriter& rewriter) {
|
|
auto enlargeTiedDpsChain = [&](Value value, RankedTensorType newType, auto& self) -> void {
|
|
auto* definingOp = value.getDefiningOp();
|
|
if (!definingOp)
|
|
return;
|
|
auto dpsDefiningOp = dyn_cast<DestinationStyleOpInterface>(definingOp);
|
|
if (!dpsDefiningOp)
|
|
return;
|
|
auto* tiedOperand = dpsDefiningOp.getTiedOpOperand(cast<OpResult>(value));
|
|
if (!tiedOperand)
|
|
return;
|
|
Value tiedValue = tiedOperand->get();
|
|
assert(tiedValue.hasOneUse() && "Tied DPS operand expected to have a single use");
|
|
tiedValue.setType(newType);
|
|
self(tiedValue, newType, self);
|
|
};
|
|
|
|
funcOp.walk([&](PimVMMOp vmmOp) {
|
|
auto outTensorOperand = vmmOp.getOutBuf();
|
|
auto resultTensor = vmmOp.getOutRes();
|
|
auto outShape = getTensorShape(outTensorOperand);
|
|
assert(isHVectorShape(outShape));
|
|
if (outShape[1] != static_cast<int64_t>(crossbarSize)) {
|
|
auto newShape = SmallVector<int64_t> {outShape[0], static_cast<int64_t>(crossbarSize)};
|
|
auto newType = RankedTensorType::get(newShape, outTensorOperand.getType().getElementType());
|
|
enlargeTiedDpsChain(outTensorOperand, newType, enlargeTiedDpsChain);
|
|
outTensorOperand.setType(newType);
|
|
resultTensor.setType(newType);
|
|
|
|
IntegerAttr zeroAttr = rewriter.getIndexAttr(0);
|
|
IntegerAttr oneAttr = rewriter.getIndexAttr(1);
|
|
IntegerAttr oldShapeZeroAttr = rewriter.getIndexAttr(outShape[0]);
|
|
IntegerAttr oldShapeOneAttr = rewriter.getIndexAttr(outShape[1]);
|
|
SmallVector<OpFoldResult> offsets = {zeroAttr, zeroAttr};
|
|
SmallVector<OpFoldResult> sizes = {oldShapeZeroAttr, oldShapeOneAttr};
|
|
SmallVector<OpFoldResult> strides = {oneAttr, oneAttr};
|
|
rewriter.setInsertionPointAfter(vmmOp);
|
|
auto sliceOp = tensor::ExtractSliceOp::create(rewriter, vmmOp.getLoc(), resultTensor, offsets, sizes, strides);
|
|
SmallPtrSet<Operation*, 2> exceptions = {vmmOp, sliceOp};
|
|
resultTensor.replaceAllUsesExcept(sliceOp.getResult(), exceptions);
|
|
}
|
|
});
|
|
}
|
|
|
|
void SpatialToPimPass::addResultBuffer(func::ReturnOp& returnOp, IRRewriter& rewriter) {
|
|
outputTensors.reserve(returnOp->getNumOperands());
|
|
rewriter.setInsertionPointToStart(returnOp->getBlock());
|
|
for (auto returnValue : returnOp->getOperands()) {
|
|
Operation* returnValueDefiningOp = returnValue.getDefiningOp();
|
|
if (returnValueDefiningOp->hasTrait<OpTrait::ConstantLike>()) {
|
|
assert(!hasWeightAlways(returnValueDefiningOp));
|
|
outputTensors.push_back(returnValue);
|
|
}
|
|
else {
|
|
auto newOutputTensor =
|
|
createEmptyTensorFromShaped(rewriter, returnValue.getLoc(), cast<ShapedType>(returnValue.getType()));
|
|
outputTensors.push_back(newOutputTensor);
|
|
}
|
|
}
|
|
}
|
|
|
|
void SpatialToPimPass::allocateAndInitializeCoreLocalVariables(func::FuncOp funcOp, IRRewriter& rewriter) {
|
|
Location loc = funcOp.getLoc();
|
|
|
|
auto insertMemCopyHostToDev = [&](auto valueToReplace, auto hostTensor, int64_t elementsOffset) {
|
|
auto tensorType = cast<ShapedType>(valueToReplace.getType());
|
|
Type elementType = tensorType.getElementType();
|
|
size_t elementByteSize = elementType.getIntOrFloatBitWidth() / 8;
|
|
rewriter.setInsertionPoint(getEarliestUserWithinBlock(valueToReplace));
|
|
|
|
auto deviceTensor = tensor::EmptyOp::create(rewriter, loc, tensorType.getShape(), elementType);
|
|
|
|
auto memCopyHostToDevOp = PimMemCopyHostToDevOp::create(
|
|
rewriter,
|
|
loc,
|
|
tensorType,
|
|
deviceTensor,
|
|
hostTensor,
|
|
rewriter.getI32IntegerAttr(0),
|
|
rewriter.getI32IntegerAttr(static_cast<int32_t>(elementsOffset * elementByteSize)),
|
|
rewriter.getI32IntegerAttr(static_cast<int32_t>(tensorType.getNumElements() * elementByteSize)));
|
|
|
|
rewriter.replaceAllUsesWith(valueToReplace, memCopyHostToDevOp.getResult());
|
|
};
|
|
|
|
// Replace input tensors with memRefs
|
|
SmallVector<bufferization::ToTensorOp, 8> inputTensors;
|
|
for (size_t i = 0; i < funcOp.getNumArguments(); i++) {
|
|
BlockArgument tensorArg = funcOp.getArgument(i);
|
|
DictionaryAttr tensorArgAttrs = funcOp.getArgAttrDict(i);
|
|
ShapedType tensorArgType = cast<ShapedType>(tensorArg.getType());
|
|
MemRefType memRefArgType = MemRefType::get(tensorArgType.getShape(), tensorArgType.getElementType());
|
|
|
|
funcOp.insertArgument(i + 1, memRefArgType, tensorArgAttrs, loc);
|
|
BlockArgument memRefArg = funcOp.getArgument(i + 1);
|
|
|
|
Block& block = funcOp.getBody().front();
|
|
rewriter.setInsertionPoint(&block.front());
|
|
auto toTensorOp =
|
|
bufferization::ToTensorOp::create(rewriter, loc, tensorArgType, memRefArg, rewriter.getUnitAttr());
|
|
inputTensors.push_back(toTensorOp);
|
|
|
|
tensorArg.replaceAllUsesWith(toTensorOp);
|
|
funcOp.eraseArgument(i);
|
|
}
|
|
|
|
llvm::SmallSet<tensor::ExtractSliceOp, 8> sliceOpsToRemove;
|
|
for (auto& op : funcOp.getBody().getOps())
|
|
if (auto computeOp = dyn_cast<spatial::SpatWeightedCompute>(op)) {
|
|
unsigned numComputeWeights = computeOp.getWeights().size();
|
|
for (auto [computeInputIdx, computeOpInput] : llvm::enumerate(computeOp.getInputs())) {
|
|
TypedValue<TensorType> tensorSource;
|
|
int64_t elementsOffset = 0;
|
|
|
|
if (auto sliceOp = dyn_cast<tensor::ExtractSliceOp>(computeOpInput.getDefiningOp())) {
|
|
tensorSource = cast<TypedValue<TensorType>>(sliceOp.getSource());
|
|
|
|
ArrayRef<int64_t> sourceShape = tensorSource.getType().getShape();
|
|
ArrayRef<int64_t> sliceOffsets = sliceOp.getStaticOffsets();
|
|
ArrayRef<int64_t> sliceSizes = sliceOp.getStaticSizes();
|
|
ArrayRef<int64_t> sliceStrides = sliceOp.getStaticStrides();
|
|
assert("Extracting slice non-contiguous in memory"
|
|
&& isMemoryContiguous(sourceShape, sliceOffsets, sliceSizes, sliceStrides));
|
|
|
|
for (size_t i = 0; i < sliceOffsets.size(); i++) {
|
|
int64_t partialOffset = sliceOffsets[i];
|
|
if (partialOffset != 0)
|
|
for (size_t j = i + 1; j < sourceShape.size(); j++)
|
|
partialOffset *= sourceShape[j];
|
|
elementsOffset += partialOffset;
|
|
}
|
|
|
|
computeOp.setOperand(numComputeWeights + computeInputIdx, tensorSource);
|
|
sliceOpsToRemove.insert(sliceOp);
|
|
}
|
|
else
|
|
tensorSource = cast<TypedValue<TensorType>>(computeOpInput);
|
|
|
|
// Compute results must be transferred through channels via send/receive
|
|
if (isa<spatial::SpatWeightedCompute>(tensorSource.getDefiningOp()))
|
|
continue;
|
|
|
|
BlockArgument computeBlockArgToReplace = computeOp.getBody().front().getArgument(computeInputIdx);
|
|
insertMemCopyHostToDev(computeBlockArgToReplace, tensorSource, elementsOffset);
|
|
}
|
|
}
|
|
|
|
for (auto sliceOp : sliceOpsToRemove)
|
|
if (sliceOp->getUses().empty())
|
|
rewriter.eraseOp(sliceOp);
|
|
}
|
|
|
|
void SpatialToPimPass::replaceBlockArgumentWithRecvOp(spatial::SpatWeightedCompute& computeOp,
|
|
unsigned int argIndex,
|
|
spatial::SpatChannelNewOp& channel,
|
|
Type& tensorType,
|
|
bool useBroadcastOp,
|
|
IRRewriter& rewriter) {
|
|
auto& computeBlock = computeOp.getRegion().front();
|
|
//(remember that WeightedCompute have weights as first operands, however these
|
|
// weights are not included in the block arguments. Thus, when indexing the
|
|
// block argument we need to remove the weights count)
|
|
auto computeWeightsCount = computeOp.getWeights().size();
|
|
auto blockArg = computeBlock.getArgument(argIndex - computeWeightsCount);
|
|
// Receive the tensor just before the first use of the value
|
|
rewriter.setInsertionPoint(getEarliestUserWithinBlock(blockArg));
|
|
Value receivedValue;
|
|
if (useBroadcastOp)
|
|
receivedValue = spatial::SpatChannelBroadcastReceiveOp::create(rewriter, computeOp.getLoc(), tensorType, channel);
|
|
else
|
|
receivedValue = spatial::SpatChannelReceiveOp::create(rewriter, computeOp.getLoc(), tensorType, channel);
|
|
|
|
blockArg.replaceAllUsesWith(receivedValue);
|
|
}
|
|
|
|
void SpatialToPimPass::addReceiveOps(Value& channelSourceOp,
|
|
spatial::SpatChannelNewOp& channel,
|
|
Type& channelTensorType,
|
|
bool& useBroadcastOp,
|
|
IRRewriter& rewriter) {
|
|
auto sourceOpUses = channelSourceOp.getUses();
|
|
|
|
// Check if we need to update `useBroadcastOp` to true, in the case of a reshapeOp with multiple users
|
|
if (useBroadcastOp == false) {
|
|
// if useBroadcastOp is false, then sourceOp must have only one user
|
|
assert(rangeLength(sourceOpUses) == 1);
|
|
|
|
if (auto reshapeOp = dyn_cast<tosa::ReshapeOp>(sourceOpUses.begin()->getOwner())) {
|
|
auto reshapeOpUses = reshapeOp.getOutput().getUses();
|
|
auto reshapeOpUsesCount = rangeLength(reshapeOpUses);
|
|
if (reshapeOpUsesCount > 1)
|
|
useBroadcastOp = true;
|
|
}
|
|
}
|
|
|
|
for (auto& resultUse : sourceOpUses) {
|
|
// The user must be a ComputeOp, or a reshapeOp which can be used by many ComputeOps
|
|
spatial::SpatWeightedCompute computeUser = dyn_cast<spatial::SpatWeightedCompute>(resultUse.getOwner());
|
|
|
|
if (computeUser) {
|
|
replaceBlockArgumentWithRecvOp(
|
|
computeUser, resultUse.getOperandNumber(), channel, channelTensorType, useBroadcastOp, rewriter);
|
|
continue;
|
|
}
|
|
|
|
if (!computeUser) {
|
|
auto reshapeOp = dyn_cast<tosa::ReshapeOp>(resultUse.getOwner());
|
|
if (!reshapeOp) {
|
|
channelSourceOp.getDefiningOp()->getParentOp()->getParentOp()->dump();
|
|
resultUse.getOwner()->dump();
|
|
llvm_unreachable("User of Value that now needs to be received by channel is not a ComputeOp nor a ReshapeOp");
|
|
}
|
|
|
|
// The tensorType now becomes the one of the reshapeOp
|
|
channelTensorType = reshapeOp.getResult().getType();
|
|
|
|
for (auto& reshapeUse : reshapeOp.getOutput().getUses()) {
|
|
computeUser = dyn_cast<spatial::SpatWeightedCompute>(reshapeUse.getOwner());
|
|
|
|
if (!computeUser)
|
|
llvm_unreachable("ReshapeOp users must be ComputeOps");
|
|
|
|
replaceBlockArgumentWithRecvOp(
|
|
computeUser, reshapeUse.getOperandNumber(), channel, channelTensorType, useBroadcastOp, rewriter);
|
|
}
|
|
|
|
// Remove the reshapeOp, so that the sourceOp has no users
|
|
operationsToRemove.push_back(reshapeOp);
|
|
}
|
|
}
|
|
}
|
|
|
|
void SpatialToPimPass::replaceReturnOpOperands(func::ReturnOp& returnOp, IRRewriter& rewriter) {
|
|
for (auto it : llvm::enumerate(returnOp.getOperands())) {
|
|
Operation* returnOperand = it.value().getDefiningOp();
|
|
|
|
size_t orderWithinReturn = it.index();
|
|
|
|
rewriter.modifyOpInPlace(returnOp,
|
|
[&] { returnOp.setOperand(orderWithinReturn, outputTensors[orderWithinReturn]); });
|
|
|
|
// If the operand is a concatenation operation and the returnOp was the only
|
|
// user of the returnOperand, we can safely remove it
|
|
if (isAConcatOp(returnOperand)) {
|
|
auto returnOperandUses = it.value().getUses();
|
|
if (rangeLength(returnOperandUses) == 0)
|
|
rewriter.eraseOp(returnOperand);
|
|
}
|
|
}
|
|
}
|
|
|
|
void SpatialToPimPass::runOnReceiveOp(spatial::SpatChannelReceiveOp receiveOp, IRRewriter& rewriter) {
|
|
|
|
auto channel = cast<spatial::SpatChannelNewOp>(receiveOp.getChannel().getDefiningOp());
|
|
|
|
auto sendOpOpt = getOtherEndOfChannel(receiveOp, true, rewriter);
|
|
if (failed(sendOpOpt))
|
|
llvm_unreachable("ChannelReceiveOp has no matching SendOp");
|
|
|
|
auto sendOp = cast<spatial::SpatChannelSendOp>(*sendOpOpt);
|
|
|
|
auto tensorType = receiveOp.getType();
|
|
Value receiveRes = receiveOp.getResult();
|
|
|
|
// Check if the receiveOp value has more than one user
|
|
auto receiveUses = receiveRes.getUses();
|
|
auto receiveUsesCount = rangeLength(receiveUses);
|
|
assert(receiveUsesCount > 0);
|
|
bool useBroadcastOp = receiveUsesCount > 1;
|
|
addReceiveOps(receiveRes, channel, tensorType, useBroadcastOp, rewriter);
|
|
|
|
if (useBroadcastOp) {
|
|
// When receiving, we actually noticed that the value has more than one
|
|
// user. This means that we need to get the replace the original SendOp with
|
|
// a BroadcastSendOp
|
|
rewriter.setInsertionPoint(sendOp);
|
|
rewriter.replaceOpWithNewOp<spatial::SpatChannelBroadcastSendOp>(sendOp, sendOp.getChannel(), sendOp.getData());
|
|
}
|
|
}
|
|
|
|
std::unique_ptr<Pass> createSpatialToPimPass() { return std::make_unique<SpatialToPimPass>(); }
|
|
|
|
} // namespace onnx_mlir
|