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

big refactor
Validate Operations / validate-operations (push) Has been cancelled
Details

Browse Source
This commit is contained in:
NiccoloN
2026-08-04 11:28:05 +02:00
parent f4a3b012cc
commit 10b6ee6c32
150 changed files with 6737 additions and 4816 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.cpp
+118 -32
View File
@@ -1,49 +1,133 @@
#include "ConvGeometry.hpp"
#include <algorithm>
#include <limits>
#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
namespace onnx_mlir {
namespace {
static const ONNXToSpatialPlanningOptions& defaultPlanningOptions() {
static const ONNXToSpatialPlanningOptions options {
std::numeric_limits<uint64_t>::max(),
std::numeric_limits<uint64_t>::max(),
spatial::ConvLoweringStrategy::Auto,
false,
};
return options;
}
} // namespace
const ONNXToSpatialPlanningOptions& ConvLoweringState::planningOptions() const {
return options ? *options : defaultPlanningOptions();
}
bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup) {
return group == numChannelsIn && numChannelsInPerGroup == 1 && numChannelsOut % group == 0;
}
ConvGeometry buildConvGeometry(const ConvLoweringState& state) {
void classifyConvProblem(ConvProblem& problem) {
problem.isDepthwise = isDepthwiseConv(
problem.group, problem.numChannelsIn, problem.numChannelsOut,
problem.numChannelsInPerGroup);
problem.isGrouped = problem.group > 1;
problem.isPointwise = problem.wHeight == 1 && problem.wWidth == 1
&& problem.strideHeight == 1 && problem.strideWidth == 1
&& problem.dilationHeight == 1 && problem.dilationWidth == 1
&& problem.padHeightBegin == 0 && problem.padHeightEnd == 0
&& problem.padWidthBegin == 0 && problem.padWidthEnd == 0;
}
ConvGeometry buildConvGeometry(const ConvProblem& problem,
const spatial::SpatialTargetResources& target) {
ConvGeometry geo {
state.batchSize,
state.numChannelsIn,
state.xHeight,
state.xWidth,
state.numChannelsOut,
state.wHeight,
state.wWidth,
state.outHeight,
state.outWidth,
state.group,
state.numChannelsInPerGroup,
state.numChannelsOutPerGroup,
state.numChannelsInPerGroup * state.wHeight * state.wWidth,
state.numChannelsOutPerGroup,
state.batchSize * state.outHeight * state.outWidth,
static_cast<int64_t>(crossbarSize.getValue()),
problem.numChannelsInPerGroup * problem.wHeight * problem.wWidth,
problem.numChannelsOutPerGroup,
problem.batchSize * problem.outHeight * problem.outWidth,
static_cast<int64_t>(target.matrixShape.rows),
static_cast<int64_t>(target.matrixUnitsPerProcessor),
1,
0,
state.hasBias,
isDepthwiseConv(state.group, state.numChannelsIn, state.numChannelsOut, state.numChannelsInPerGroup),
};
geo.pack = std::max<int64_t>(1, geo.xbarSize / std::max<int64_t>(geo.k, geo.c));
geo.im2colElements = static_cast<uint64_t>(std::max<int64_t>(0, geo.p)) * static_cast<uint64_t>(std::max<int64_t>(0, geo.k));
return geo;
}
uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor) {
static ConvMaterializationKind getMaterializationKind(
spatial::ConvLoweringStrategy strategy) {
switch (strategy) {
case spatial::ConvLoweringStrategy::Depthwise:
return ConvMaterializationKind::StructuredDepthwise;
case spatial::ConvLoweringStrategy::Legacy:
case spatial::ConvLoweringStrategy::PackedIm2Col:
return ConvMaterializationKind::PackedIm2Col;
case spatial::ConvLoweringStrategy::StreamedPatch:
case spatial::ConvLoweringStrategy::OutputChannelTiled:
case spatial::ConvLoweringStrategy::Tiled2D:
return ConvMaterializationKind::StreamedPatch;
case spatial::ConvLoweringStrategy::StreamedPacked:
return ConvMaterializationKind::StreamedPacked;
case spatial::ConvLoweringStrategy::InputKTiled:
return ConvMaterializationKind::InputKTiled;
case spatial::ConvLoweringStrategy::Auto:
break;
}
llvm_unreachable("auto is not a Conv materialization kind");
}
static bool fitsSingleCrossbar(const ConvGeometry& geo) {
return geo.k <= geo.xbarSize && geo.c <= geo.xbarSize;
}
static bool fitsPackedIm2Col(const ConvGeometry& geo,
const ONNXToSpatialPlanningOptions& options) {
return fitsSingleCrossbar(geo) && geo.pack >= 2
&& geo.im2colElements <= options.convIm2colMaxElements;
}
mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
spatial::ConvLoweringStrategy strategy,
const spatial::SpatialTargetResources& target,
const ONNXToSpatialPlanningOptions& options) {
ConvGeometry geo = buildConvGeometry(problem, target);
auto plan = [&]() { return ConvPlan {getMaterializationKind(strategy)}; };
auto ifApplicable = [&](bool applicable) -> mlir::FailureOr<ConvPlan> {
return applicable ? mlir::FailureOr<ConvPlan>(plan()) : mlir::FailureOr<ConvPlan>(mlir::failure());
};
switch (strategy) {
case spatial::ConvLoweringStrategy::Auto:
return mlir::failure();
case spatial::ConvLoweringStrategy::Legacy:
return plan();
case spatial::ConvLoweringStrategy::Depthwise:
return ifApplicable(problem.isDepthwise);
case spatial::ConvLoweringStrategy::PackedIm2Col:
return ifApplicable(fitsPackedIm2Col(geo, options));
case spatial::ConvLoweringStrategy::StreamedPatch:
return ifApplicable(fitsSingleCrossbar(geo));
case spatial::ConvLoweringStrategy::StreamedPacked:
return ifApplicable(fitsSingleCrossbar(geo) && geo.pack >= 2);
case spatial::ConvLoweringStrategy::OutputChannelTiled:
return ifApplicable(geo.k <= geo.xbarSize && geo.c > geo.xbarSize);
case spatial::ConvLoweringStrategy::InputKTiled:
return ifApplicable(geo.k > geo.xbarSize && geo.c <= geo.xbarSize);
case spatial::ConvLoweringStrategy::Tiled2D:
return ifApplicable(geo.k > geo.xbarSize && geo.c > geo.xbarSize);
}
llvm_unreachable("unknown Conv lowering strategy");
}
uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
int64_t packFactor,
const ONNXToSpatialPlanningOptions& options) {
const uint64_t patchElements = static_cast<uint64_t>(std::max<int64_t>(1, geo.k));
uint64_t chunkPositions = std::max<uint64_t>(1, pimConvIm2colMaxElements / patchElements);
uint64_t chunkPositions = std::max<uint64_t>(1, options.convIm2colMaxElements / patchElements);
chunkPositions = std::min<uint64_t>(chunkPositions, static_cast<uint64_t>(std::max<int64_t>(1, geo.p)));
chunkPositions = std::min<uint64_t>(chunkPositions, std::max<uint64_t>(1, pimConvStreamChunkPositions));
chunkPositions = std::min<uint64_t>(chunkPositions, std::max<uint64_t>(1, options.convStreamChunkPositions));
if (packFactor > 1 && chunkPositions > static_cast<uint64_t>(packFactor)) {
chunkPositions -= chunkPositions % static_cast<uint64_t>(packFactor);
@@ -52,24 +136,26 @@ uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor)
return std::max<uint64_t>(1, chunkPositions);
}
RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvLoweringState& state) {
const int64_t rawBegin = outputRows.begin * state.strideHeight - state.padHeightBegin;
RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvProblem& problem) {
const int64_t rawBegin = outputRows.begin * problem.strideHeight - problem.padHeightBegin;
const int64_t rawEnd =
(outputRows.end - 1) * state.strideHeight - state.padHeightBegin + state.dilationHeight * (state.wHeight - 1) + 1;
return {std::max<int64_t>(0, rawBegin), std::min<int64_t>(state.xHeight, rawEnd)};
(outputRows.end - 1) * problem.strideHeight - problem.padHeightBegin
+ problem.dilationHeight * (problem.wHeight - 1) + 1;
return {std::max<int64_t>(0, rawBegin), std::min<int64_t>(problem.xHeight, rawEnd)};
}
ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvLoweringState& state) {
ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvProblem& problem) {
ConvRowDemand demand;
demand.outputRows = outputRows;
demand.neededInputRows = computeConvInputRowsForOutputRows(outputRows, state);
demand.neededInputRows = computeConvInputRowsForOutputRows(outputRows, problem);
demand.acquiredInputRows = demand.neededInputRows;
const int64_t rawBegin = outputRows.begin * state.strideHeight - state.padHeightBegin;
const int64_t rawBegin = outputRows.begin * problem.strideHeight - problem.padHeightBegin;
const int64_t rawEnd =
(outputRows.end - 1) * state.strideHeight - state.padHeightBegin + state.dilationHeight * (state.wHeight - 1) + 1;
(outputRows.end - 1) * problem.strideHeight - problem.padHeightBegin
+ problem.dilationHeight * (problem.wHeight - 1) + 1;
demand.topHaloRows = std::max<int64_t>(0, -rawBegin);
demand.bottomHaloRows = std::max<int64_t>(0, rawEnd - state.xHeight);
demand.bottomHaloRows = std::max<int64_t>(0, rawEnd - problem.xHeight);
demand.acquiredInputRows = demand.neededInputRows;
return demand;
}