replace deprecated "rewriter.create()" pattern
refactor PIM to Pim everywhere except for the accelerator name
This commit is contained in:
NiccoloN
551
src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp
Normal file
551
src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp
Normal file
@@ -0,0 +1,551 @@
|
||||
#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
|
||||
Reference in New Issue
Block a user