nnicolosi/Raptor
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add .clang-format

Browse Source 
reformat all src
This commit is contained in:
NiccoloN
2026-02-26 19:16:42 +01:00
parent a2c31836ae
commit 810e5e75f9
32 changed files with 902 additions and 953 deletions
Download Patch File Download Diff File Expand all files Collapse all files

194
src/PIM/Conversion/ONNXToSpatial/Utils/SpatialReducer.cpp
View File

@@ -1,32 +1,31 @@
#include "SpatialReducer.hpp"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/Value.h"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "llvm/Support/raw_ostream.h"
#include <cassert>
#include <unordered_map>
#include <utility>
#include "SpatialReducer.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#define GET_COMP(computeOpAndResNum) std::get<0>(computeOpAndResNum)
#define GET_RES_NUM(computeOpAndResNum) std::get<1>(computeOpAndResNum)
namespace onnx_mlir {
llvm::SmallPtrSet<Operation *, 16>
onnx_mlir::SpatialReducer::oldComputeOpsReplaced;
llvm::SmallPtrSet<Operation*, 16> onnx_mlir::SpatialReducer::oldComputeOpsReplaced;
ResNum SpatialReducer::applyResultProcessing(
ComputeAndResNum computeOpAndResNum,
std::function<Value(const Value &)> processFun,
ConversionPatternRewriter &rewriter) {
ResNum SpatialReducer::applyResultProcessing(ComputeAndResNum computeOpAndResNum,
std::function<Value(const Value&)> processFun,
ConversionPatternRewriter& rewriter) {
assert(processFun);
auto computeOp = GET_COMP(computeOpAndResNum);
auto resultNum = GET_RES_NUM(computeOpAndResNum);
spatial::SpatYieldOp yieldOp =
cast<spatial::SpatYieldOp>(computeOp.getBody().front().getTerminator());
spatial::SpatYieldOp yieldOp = cast<spatial::SpatYieldOp>(computeOp.getBody().front().getTerminator());
Value result = yieldOp->getOperand(resultNum);
rewriter.setInsertionPointAfterValue(result);
@@ -43,30 +42,24 @@ ResNum SpatialReducer::applyResultProcessing(
return yieldOp.getNumOperands() - 1;
}
OpAndResNum SpatialReducer::applyReducePattern(
SmallVector<ComputeAndResNum> &computeOpsAndResNum,
std::function<Value(const Value &, const Value &)> reduce,
std::function<Value(const Value &)> preprocess,
std::function<Value(const Value &)> postprocess) {
OpAndResNum SpatialReducer::applyReducePattern(SmallVector<ComputeAndResNum>& computeOpsAndResNum,
std::function<Value(const Value&, const Value&)> reduce,
std::function<Value(const Value&)> preprocess,
std::function<Value(const Value&)> postprocess) {
if (preprocess) {
for (auto &computeOpAndResNum : computeOpsAndResNum) {
GET_RES_NUM(computeOpAndResNum) =
applyResultProcessing(computeOpAndResNum, preprocess, rewriter);
}
}
if (preprocess)
for (auto& computeOpAndResNum : computeOpsAndResNum)
GET_RES_NUM(computeOpAndResNum) = applyResultProcessing(computeOpAndResNum, preprocess, rewriter);
// It is possible that `computeOpsAndResNum` contains two entries for the same
// computeOp. In this case, we need to apply the reduction within-computef
// Keep a map between a computeOp and the last Value for this reduction
std::unordered_map<Operation *, Value> lastValueForCompute;
for (auto &computeOpAndResNum : computeOpsAndResNum) {
std::unordered_map<Operation*, Value> lastValueForCompute;
for (auto& computeOpAndResNum : computeOpsAndResNum) {
auto computeOp = GET_COMP(computeOpAndResNum);
auto yieldOp =
cast<spatial::SpatYieldOp>(computeOp.getBody().front().getTerminator());
Value valueWithinCompute =
yieldOp->getOperand(GET_RES_NUM(computeOpAndResNum));
auto yieldOp = cast<spatial::SpatYieldOp>(computeOp.getBody().front().getTerminator());
Value valueWithinCompute = yieldOp->getOperand(GET_RES_NUM(computeOpAndResNum));
auto it = lastValueForCompute.find(computeOp.getOperation());
@@ -75,15 +68,12 @@ OpAndResNum SpatialReducer::applyReducePattern(
// within-compute
Value lastWithinComputeValue = it->second;
assert(valueWithinCompute.getDefiningOp() &&
lastWithinComputeValue.getDefiningOp());
assert(valueWithinCompute.getDefiningOp() && lastWithinComputeValue.getDefiningOp());
if (valueWithinCompute.getDefiningOp()->isBeforeInBlock(
lastWithinComputeValue.getDefiningOp())) {
if (valueWithinCompute.getDefiningOp()->isBeforeInBlock(lastWithinComputeValue.getDefiningOp()))
rewriter.setInsertionPointAfterValue(lastWithinComputeValue);
} else {
else
rewriter.setInsertionPointAfterValue(valueWithinCompute);
}
valueWithinCompute = reduce(lastWithinComputeValue, valueWithinCompute);
lastValueForCompute[computeOp.getOperation()] = valueWithinCompute;
}
@@ -94,16 +84,15 @@ OpAndResNum SpatialReducer::applyReducePattern(
// Now, reconstruct from the map the computeOpsAndResNum list
computeOpsAndResNum.clear();
computeOpsAndResNum.reserve(lastValueForCompute.size());
for (auto &entry : lastValueForCompute) {
for (auto& entry : lastValueForCompute) {
auto computeOp = cast<spatial::SpatWeightedCompute>(entry.first);
auto valueWithinCompute = entry.second;
// We check if `valueWithinCompute` is already used by the yieldOp, in that
// case no need to add it
auto yieldOp =
cast<spatial::SpatYieldOp>(computeOp.getBody().front().getTerminator());
auto yieldOp = cast<spatial::SpatYieldOp>(computeOp.getBody().front().getTerminator());
bool yieldOpUseFound = false;
for (auto &use : valueWithinCompute.getUses()) {
for (auto& use : valueWithinCompute.getUses()) {
if (use.getOwner() == yieldOp.getOperation()) {
// If the value is already used by the yieldOp, we can just use it
computeOpsAndResNum.push_back({computeOp, use.getOperandNumber()});
@@ -111,9 +100,8 @@ OpAndResNum SpatialReducer::applyReducePattern(
break;
}
}
if (yieldOpUseFound) {
if (yieldOpUseFound)
continue;
}
// If this result is not used within a yieldOp, then add it
auto resultNum = yieldOp->getNumOperands();
@@ -147,23 +135,18 @@ OpAndResNum SpatialReducer::applyReducePattern(
// the number of results)
// See below `reducerChanges.push_back` and `finalizeReduceUpdates`
auto yieldOpFirstCompute = cast<spatial::SpatYieldOp>(
firstCompute.getBody().front().getTerminator());
auto yieldOpFirstCompute = cast<spatial::SpatYieldOp>(firstCompute.getBody().front().getTerminator());
// Add a new operand to the block of the second computeOp
Block &secondBlock = secondCompute.getBody().front();
Value formerRes1 = secondBlock.addArgument(
yieldOpFirstCompute->getOperand(firstResultNum).getType(), loc);
Block& secondBlock = secondCompute.getBody().front();
Value formerRes1 = secondBlock.addArgument(yieldOpFirstCompute->getOperand(firstResultNum).getType(), loc);
auto secondComputeWeightsNum =
secondCompute->getAttrOfType<DenseI32ArrayAttr>(
secondCompute.getOperandSegmentSizesAttrName())[0];
auto secondComputeOperandNum =
secondComputeWeightsNum + secondBlock.getNumArguments() - 1;
secondCompute->getAttrOfType<DenseI32ArrayAttr>(secondCompute.getOperandSegmentSizesAttrName())[0];
auto secondComputeOperandNum = secondComputeWeightsNum + secondBlock.getNumArguments() - 1;
// Take the "former-result" from the second computeOp
spatial::SpatYieldOp secondYield =
cast<spatial::SpatYieldOp>(secondBlock.getTerminator());
spatial::SpatYieldOp secondYield = cast<spatial::SpatYieldOp>(secondBlock.getTerminator());
Value formerRes2 = secondYield.getOperand(secondResultNum);
// Apply reduction operation
@@ -184,37 +167,31 @@ OpAndResNum SpatialReducer::applyReducePattern(
// We should also add an entry for updating the results of the last
// operation (the one which never becomes a `firstCompute`): because it is
// not tracked by reducerChanges as `fromOp`
reducerChanges.push_back({firstCompute.getOperation(), firstResultNum,
secondCompute.getOperation(), secondComputeOperandNum});
reducerChanges.push_back(
{firstCompute.getOperation(), firstResultNum, secondCompute.getOperation(), secondComputeOperandNum});
nextComputeOps.push_back(std::make_pair(secondCompute, secondResultNum));
}
// If we have an odd number of inputs, we need to add the last one to the
// newInputs list.
if (computeOpsRef.size() % 2 == 1) {
if (computeOpsRef.size() % 2 == 1)
nextComputeOps.push_back(computeOpsRef.back());
}
// Replace the inputOps list with the new one.
computeOpsRef =
llvm::OwningArrayRef<ComputeAndResNum>(std::move(nextComputeOps));
computeOpsRef = llvm::OwningArrayRef<ComputeAndResNum>(std::move(nextComputeOps));
}
assert(computeOpsRef.size() == 1 &&
"Internal error: expected a single input at this point.");
assert(computeOpsRef.size() == 1 && "Internal error: expected a single input at this point.");
auto finalComputeAndResNum = computeOpsRef[0];
// Force the update of the results of this computeOp, when finalizing
computeOpNeedingResUpdate.push_back(GET_COMP(finalComputeAndResNum));
if (postprocess) {
GET_RES_NUM(finalComputeAndResNum) =
applyResultProcessing(finalComputeAndResNum, postprocess, rewriter);
}
if (postprocess)
GET_RES_NUM(finalComputeAndResNum) = applyResultProcessing(finalComputeAndResNum, postprocess, rewriter);
return std::make_pair(GET_COMP(finalComputeAndResNum).getOperation(),
GET_RES_NUM(finalComputeAndResNum));
return std::make_pair(GET_COMP(finalComputeAndResNum).getOperation(), GET_RES_NUM(finalComputeAndResNum));
}
void SpatialReducer::finalizeReduceUpdates() {
@@ -223,15 +200,13 @@ void SpatialReducer::finalizeReduceUpdates() {
reducesFinalized = true;
// First, add the results to the computeOps
for (auto &reduceChange : reducerChanges) {
for (auto& reduceChange : reducerChanges)
updateResultsOfCompute(reduceChange.fromOp);
}
for (auto &c : computeOpNeedingResUpdate) {
for (auto& c : computeOpNeedingResUpdate)
updateResultsOfCompute(c.getOperation());
}
for (auto &reducerChange : this->reducerChanges) {
for (auto& reducerChange : this->reducerChanges) {
auto fromOp = reducerChange.fromOp;
auto toOp = reducerChange.toOp;
auto fromOpResNum = reducerChange.fromOpResNum;
@@ -243,16 +218,14 @@ void SpatialReducer::finalizeReduceUpdates() {
// toComputeOp could be the existing pointer, or we have to remap it with
// `opToReplacedCompute`
auto toComputeOp = opToReplacedCompute[toOp];
if (!toComputeOp) {
if (!toComputeOp)
toComputeOp = cast<spatial::SpatWeightedCompute>(toOp);
}
assert(toComputeOp != fromComputeOp &&
"Oops should have caught this earlier!");
assert(toComputeOp != fromComputeOp && "Oops should have caught this earlier!");
assert(toComputeOp->getNumOperands() == toOpOperandNum &&
"toOpOperandNum should be the last operand of toComputeOp, are the "
"operations in the right order?");
assert(toComputeOp->getNumOperands() == toOpOperandNum
&& "toOpOperandNum should be the last operand of toComputeOp, are the "
"operations in the right order?");
// Add the new operand to `toComputeOp`
auto fromResult = fromComputeOp.getResult(fromOpResNum);
@@ -261,24 +234,22 @@ void SpatialReducer::finalizeReduceUpdates() {
}
}
Value SpatialReducer::resolveValueFromOpAndResNum(OpAndResNum &opAndResNum) {
assert(reducesFinalized &&
"Cannot create resolve values before finalizing the reduce updates.");
Value SpatialReducer::resolveValueFromOpAndResNum(OpAndResNum& opAndResNum) {
assert(reducesFinalized && "Cannot create resolve values before finalizing the reduce updates.");
Operation *opToCast;
Operation* opToCast;
auto it = opToReplacedCompute.find(opAndResNum.first);
if (it != opToReplacedCompute.end()) {
if (it != opToReplacedCompute.end())
opToCast = it->second;
} else {
else
opToCast = opAndResNum.first;
}
auto computeOp = cast<spatial::SpatWeightedCompute>(opToCast);
return computeOp.getResult(opAndResNum.second);
}
void SpatialReducer::updateResultsOfCompute(Operation *computeOp) {
void SpatialReducer::updateResultsOfCompute(Operation* computeOp) {
if (opToReplacedCompute.find(computeOp) != opToReplacedCompute.end()) {
// If we have already replaced the fromOp, we do not need to do it again
return;
@@ -287,8 +258,7 @@ void SpatialReducer::updateResultsOfCompute(Operation *computeOp) {
auto oldComputeOpNum = oldComputeOp->getNumOperands();
auto yieldOp =
cast<spatial::SpatYieldOp>(oldComputeOp.getBody().front().getTerminator());
auto yieldOp = cast<spatial::SpatYieldOp>(oldComputeOp.getBody().front().getTerminator());
if (yieldOp.getNumOperands() == oldComputeOp->getNumResults()) {
// No result was added, just add itself to the map
@@ -301,9 +271,8 @@ void SpatialReducer::updateResultsOfCompute(Operation *computeOp) {
// Create a new ComputeOp with the new result type, but same operands
rewriter.setInsertionPoint(oldComputeOp);
auto newComputeOp =
rewriter.create<spatial::SpatWeightedCompute>(oldComputeOp->getLoc(),
newResultTypes, oldComputeOp.getWeights(), oldComputeOp.getInputs());
auto newComputeOp = rewriter.create<spatial::SpatWeightedCompute>(
oldComputeOp->getLoc(), newResultTypes, oldComputeOp.getWeights(), oldComputeOp.getInputs());
newComputeOp.getBody().takeBody(oldComputeOp.getBody());
@@ -329,54 +298,49 @@ void SpatialReducer::updateResultsOfCompute(Operation *computeOp) {
rewriter.eraseOp(oldComputeOp);
}
Value SpatialReducer::createImgConcatOp(
SmallVector<SmallVector<SmallVector<OpAndResNum>>> &outputTiles,
Location &loc, Type outputType) {
Value SpatialReducer::createImgConcatOp(SmallVector<SmallVector<SmallVector<OpAndResNum>>>& outputTiles,
Location& loc,
Type outputType) {
assert(reducesFinalized &&
"Cannot create ImgConcatOp before finalizing the reduce updates.");
assert(reducesFinalized && "Cannot create ImgConcatOp before finalizing the reduce updates.");
// outputTiles are indexed like this: [channelTile][x][y]
auto tilesCount = outputTiles.size();
auto width = outputTiles[0].size();
auto height = outputTiles[0][0].size();
SmallVector<SmallVector<SmallVector<Value>>> remappedOutputTiles(tilesCount,
SmallVector<SmallVector<Value>>(width, SmallVector<Value>(height)));
SmallVector<SmallVector<SmallVector<Value>>> remappedOutputTiles(
tilesCount, SmallVector<SmallVector<Value>>(width, SmallVector<Value>(height)));
for (size_t t = 0; t < tilesCount; t++)
for (size_t x = 0; x < width; x++)
for (size_t y = 0; y < height; y++)
remappedOutputTiles[t][x][y] =
resolveValueFromOpAndResNum(outputTiles[t][x][y]);
remappedOutputTiles[t][x][y] = resolveValueFromOpAndResNum(outputTiles[t][x][y]);
return ::onnx_mlir::createImgConcatOp(
remappedOutputTiles, rewriter, loc, outputType);
return ::onnx_mlir::createImgConcatOp(remappedOutputTiles, rewriter, loc, outputType);
}
OpAndResNum SpatialReducer::applyAddMapReduction(
SmallVector<ComputeAndResNum> &computeOps,
ConversionPatternRewriter &rewriter, Value biasTile, MapOperations mapOp) {
OpAndResNum SpatialReducer::applyAddMapReduction(SmallVector<ComputeAndResNum>& computeOps,
ConversionPatternRewriter& rewriter,
Value biasTile,
MapOperations mapOp) {
std::function<Value(const Value &)> postprocessing = nullptr;
std::function<Value(const Value&)> postprocessing = nullptr;
if (mapOp != MapOperations::None) {
postprocessing = [&](const Value a) {
Value mapOperand = a;
if (biasTile) {
mapOperand = rewriter.create<spatial::SpatVAddOp>(
a.getLoc(), a.getType(), a, biasTile);
}
if (biasTile)
mapOperand = rewriter.create<spatial::SpatVAddOp>(a.getLoc(), a.getType(), a, biasTile);
return createMapOperation(rewriter, mapOp, mapOperand);
};
}
return this->applyReducePattern(
computeOps,
[&](Value a, Value b) {
return rewriter.create<spatial::SpatVAddOp>(a.getLoc(), a.getType(), a, b);
},
/* preprocess = */ nullptr, postprocessing);
computeOps,
[&](Value a, Value b) { return rewriter.create<spatial::SpatVAddOp>(a.getLoc(), a.getType(), a, b); },
/* preprocess = */ nullptr,
postprocessing);
}
} // namespace onnx_mlir
} // namespace onnx_mlir