NiccoloN
78 lines
3.2 KiB
C++
78 lines
3.2 KiB
C++
#include "ConvGeometry.hpp"
|
|
|
|
#include <algorithm>
|
|
|
|
#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
|
|
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
|
|
|
|
namespace onnx_mlir {
|
|
|
|
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) {
|
|
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()),
|
|
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) {
|
|
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);
|
|
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));
|
|
|
|
if (packFactor > 1 && chunkPositions > static_cast<uint64_t>(packFactor)) {
|
|
chunkPositions -= chunkPositions % static_cast<uint64_t>(packFactor);
|
|
chunkPositions = std::max<uint64_t>(chunkPositions, static_cast<uint64_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;
|
|
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)};
|
|
}
|
|
|
|
ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvLoweringState& state) {
|
|
ConvRowDemand demand;
|
|
demand.outputRows = outputRows;
|
|
demand.neededInputRows = computeConvInputRowsForOutputRows(outputRows, state);
|
|
demand.acquiredInputRows = demand.neededInputRows;
|
|
|
|
const int64_t rawBegin = outputRows.begin * state.strideHeight - state.padHeightBegin;
|
|
const int64_t rawEnd =
|
|
(outputRows.end - 1) * state.strideHeight - state.padHeightBegin + state.dilationHeight * (state.wHeight - 1) + 1;
|
|
demand.topHaloRows = std::max<int64_t>(0, -rawBegin);
|
|
demand.bottomHaloRows = std::max<int64_t>(0, rawEnd - state.xHeight);
|
|
demand.acquiredInputRows = demand.neededInputRows;
|
|
return demand;
|
|
}
|
|
|
|
} // namespace onnx_mlir
|