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

remove old unused stuff

Browse Source
This commit is contained in:
NiccoloN
2026-03-23 20:00:09 +01:00
parent f2d593f749
commit f869925b64
12 changed files with 16 additions and 906 deletions
Download Patch File Download Diff File Expand all files Collapse all files

154
src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.cpp
View File

@@ -5,14 +5,12 @@
#include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/DialectConversion.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include <cassert>
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Utils/SpatialReducer.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -21,12 +19,8 @@ using namespace mlir;
namespace onnx_mlir {
namespace {
constexpr StringRef COMPUTE_HAS_SOFTMAX_DIVISOR_ATTRNAME = "computeWithSoftmaxDivisor";
static FailureOr<Value> materializeScaledConstantTensor(Value value,
float factor,
ConversionPatternRewriter& rewriter,
Location loc) {
static FailureOr<Value>
materializeScaledConstantTensor(Value value, float factor, ConversionPatternRewriter& rewriter, Location loc) {
if (factor == 1.0f)
return value;
@@ -70,16 +64,6 @@ struct GemvToSpatialCompute : OpConversionPattern<ONNXGemmOp> {
LogicalResult matchAndRewrite(ONNXGemmOp gemmOp,
ONNXGemmOpAdaptor gemmOpAdaptor,
ConversionPatternRewriter& rewriter) const override;
private:
static Value resolveONNXExpOpFromUseChain(Value startValue);
static LogicalResult softmaxReductionApplication(SmallVector<OpAndResNum>& outputOpsAndResNums,
Value& softmaxChannel,
ConversionPatternRewriter& rewriter,
SpatialReducer& reducer,
ONNXGemmOp& gemmOp,
Location& loc);
};
} // namespace
@@ -122,7 +106,13 @@ LogicalResult GemmToManyGemv::matchAndRewrite(ONNXGemmOp gemmOp,
// Expand rank-1 bias [N] to rank-2 [1, N] for uniform handling
if (cType.getRank() == 1) {
auto expandedType = RankedTensorType::get({1, cType.getDimSize(0)}, cType.getElementType());
c = tensor::ExpandShapeOp::create(rewriter, loc, expandedType, c, SmallVector<ReassociationIndices>{{0, 1}});
c = tensor::ExpandShapeOp::create(rewriter,
loc,
expandedType,
c,
SmallVector<ReassociationIndices> {
{0, 1}
});
cType = expandedType;
}
assert("Only support rank 2 tensor for C" && cType.getRank() == 2);
@@ -208,7 +198,13 @@ LogicalResult GemvToSpatialCompute::matchAndRewrite(ONNXGemmOp gemmOp,
// Expand rank-1 bias [N] to rank-2 [1, N] for uniform handling
if (cType.getRank() == 1) {
auto expandedType = RankedTensorType::get({1, cType.getDimSize(0)}, cType.getElementType());
c = tensor::ExpandShapeOp::create(rewriter, gemmLoc, expandedType, c, SmallVector<ReassociationIndices>{{0, 1}});
c = tensor::ExpandShapeOp::create(rewriter,
gemmLoc,
expandedType,
c,
SmallVector<ReassociationIndices> {
{0, 1}
});
cType = expandedType;
}
assert("Only support rank 2 tensor for C" && cType.getRank() == 2);
@@ -356,124 +352,6 @@ LogicalResult GemvToSpatialCompute::matchAndRewrite(ONNXGemmOp gemmOp,
return success();
}
Value GemvToSpatialCompute::resolveONNXExpOpFromUseChain(Value startValue) {
Value walker = startValue;
while (!llvm::isa<ONNXExpOp>(walker.getDefiningOp())) {
walker = walker.getDefiningOp()->getOperand(0);
assert(walker && walker.getDefiningOp()
&& "Unwinded the whole chain of operations while trying to "
"find ONNXExpOp, but did not find it");
}
// Make sure the dividend is actually produced by an ONNXExpOp
assert(llvm::isa<ONNXExpOp>(walker.getDefiningOp())
&& "Old output tile (softmax reducer) is not produced by an "
"ONNXExpOp");
return walker;
}
LogicalResult GemvToSpatialCompute::softmaxReductionApplication(SmallVector<OpAndResNum>& outputOpsAndResNums,
Value& softmaxChannel,
ConversionPatternRewriter& rewriter,
SpatialReducer& reducer,
ONNXGemmOp& gemmOp,
Location& loc) {
// TODO: Check case with one compute op
// Cast vector of Value into vector of ComputeOp
SmallVector<ComputeAndResNum> softmaxOpsToReduce =
llvm::to_vector(llvm::map_range(outputOpsAndResNums, [&](OpAndResNum computeAndResNum) {
return std::make_pair(cast<spatial::SpatWeightedCompute>(computeAndResNum.first), computeAndResNum.second);
}));
RankedTensorType::Builder tensorTypeBuilder({1}, Float32Type::get(rewriter.getContext()), nullptr);
const TensorType scalarTensorType = tensorTypeBuilder;
reducer.applyReducePattern(
softmaxOpsToReduce,
[&](Value a, Value b) { return spatial::SpatVAddOp::create(rewriter, loc, scalarTensorType, a, b); },
/* preprocess = */
[&](Value a) { return spatial::SpatSumOp::create(rewriter, loc, scalarTensorType, a); },
[&](Value softmaxDivisor) {
// Signal that this is the compute with the softmax divisor
auto computeOp = cast<spatial::SpatWeightedCompute>(softmaxDivisor.getDefiningOp()->getParentOp());
computeOp->setAttr(COMPUTE_HAS_SOFTMAX_DIVISOR_ATTRNAME, rewriter.getUnitAttr());
// Broadcast the divisor to all the cores
rewriter.setInsertionPointAfterValue(softmaxDivisor);
spatial::SpatChannelBroadcastSendOp::create(rewriter, loc, softmaxChannel, softmaxDivisor);
/*
* softmaxDividend = onnx.exp (...)
* sum = spat.SumOp(softmaxDividend)
* [following can be repeated N times, thus walk the use chain]
* softmaxDivisor = spat.sadd(sum, ...)
*/
Value softmaxDividend = resolveONNXExpOpFromUseChain(softmaxDivisor.getDefiningOp()->getOperand(0));
// Make sure the dividend is actually produced by an ONNXExpOp
assert(llvm::isa<ONNXExpOp>(softmaxDividend.getDefiningOp())
&& "Dividend of softmax reduction is not an ONNXExpOp");
// Do not divide here, divide after this
return softmaxDivisor;
});
// In all the cores, insert a ChannelRecvOp and divide the output tile by
// the reduced denominator.
outputOpsAndResNums.clear();
outputOpsAndResNums.reserve(softmaxOpsToReduce.size());
for (auto& computeToDivideOpAndResNum : softmaxOpsToReduce) {
auto yieldOp = cast<spatial::SpatYieldOp>(computeToDivideOpAndResNum.first.getBody().front().getTerminator());
Value divisor;
// Check if this compute contains the softmax divisor: if so, find the
// ChannelBroadcastSendOp, otherwise receive the value from the channel
// using ChannelBroadcastReceiveOp
if (computeToDivideOpAndResNum.first->hasAttr(COMPUTE_HAS_SOFTMAX_DIVISOR_ATTRNAME)) {
bool found = false;
for (auto broadcastOp :
computeToDivideOpAndResNum.first.getBody().front().getOps<spatial::SpatChannelBroadcastSendOp>()) {
assert(found == false
&& "More than one ChannelBroadcastSendOp in "
"compute? How is this possible?");
found = true;
divisor = broadcastOp.getData();
}
assert(found
&& "No ChannelBroadcastSendOp in compute where softmax "
"divisor was specified to be?");
}
else {
rewriter.setInsertionPoint(yieldOp);
divisor = spatial::SpatChannelBroadcastReceiveOp::create(rewriter, loc, scalarTensorType, softmaxChannel);
}
// Walk the chain of operations until we find the ONNXExpOp: this is
// needed because some some may have a different amount of `VAddOp`s due
// to the tree reduction (e.g. some may have no VAddOp, some may have
// multiples)
Value oldOutputTile = resolveONNXExpOpFromUseChain(yieldOp->getOperand(computeToDivideOpAndResNum.second));
rewriter.setInsertionPoint(yieldOp);
Value newOutputTile = spatial::SpatVSDivOp::create(rewriter, loc, oldOutputTile.getType(), oldOutputTile, divisor);
auto yieldOperandNum = yieldOp->getNumOperands();
yieldOp->insertOperands(yieldOperandNum, newOutputTile);
outputOpsAndResNums.push_back({computeToDivideOpAndResNum.first, yieldOperandNum});
}
return success();
}
void populateOnnxGemmOpPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
patterns.insert<GemmToManyGemv>(ctx);
patterns.insert<GemvToSpatialCompute>(ctx);