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

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

Browse Source 
slightly faster codegen
This commit is contained in:
NiccoloN
2026-07-20 16:01:57 +02:00
parent ab54243fda
commit dbb66be93e
27 changed files with 1358 additions and 1241 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp
+31 -14
View File
@@ -2180,39 +2180,55 @@ static Value createIm2colRows(const ConvLoweringState& state,
auto elemType = preparedInput.type.getElementType();
auto packedRowType = RankedTensorType::get(
{plan.effectiveMaxParallelPixels * plan.patchSize}, elemType, plan.gemmInputRowsType.getEncoding());
auto zeroAttr = DenseElementsAttr::get(packedRowType, rewriter.getZeroAttr(elemType));
Value zeroRow = getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, packedRowType);
auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
auto patchRowType = RankedTensorType::get({plan.patchSize}, elemType);
bool hasPartialLane = plan.chunkNumPatches % plan.effectiveMaxParallelPixels != 0;
SmallVector<Value> im2colInputs {preparedInput.value};
auto im2colComputeOp = createSpatComputeBatch(
rewriter,
loc,
TypeRange {plan.gemmInputRowsType},
plan.packedNumRows,
{},
ValueRange {preparedInput.value, zeroRow},
im2colInputs,
[&](detail::SpatComputeBatchBodyArgs args) {
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
Value cPack = getOrCreateIndexConstant(rewriter, anchorOp, plan.effectiveMaxParallelPixels);
Value cNumPatches = getOrCreateIndexConstant(rewriter, anchorOp, plan.chunkNumPatches);
Value laneStart = affineMulConst(rewriter, loc, args.lane, plan.effectiveMaxParallelPixels, anchorOp);
Value remaining = arith::SubIOp::create(rewriter, loc, cNumPatches, laneStart);
Value isPartial = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::ult, remaining, cPack);
Value lanePatches = arith::SelectOp::create(rewriter, loc, isPartial, remaining, cPack);
auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
auto patchRowType = RankedTensorType::get({plan.patchSize}, elemType);
Value lanePatches = cPack;
if (hasPartialLane) {
Value cNumPatches = getOrCreateIndexConstant(rewriter, anchorOp, plan.chunkNumPatches);
Value remaining = arith::SubIOp::create(rewriter, loc, cNumPatches, laneStart);
Value isPartial = arith::CmpIOp::create(
rewriter, loc, arith::CmpIPredicate::ult, remaining, cPack);
lanePatches = arith::SelectOp::create(rewriter, loc, isPartial, remaining, cPack);
}
Value rowInit = tensor::EmptyOp::create(rewriter, loc, packedRowType.getShape(), elemType);
if (hasPartialLane) {
auto zeroAttr = cast<TypedAttr>(rewriter.getZeroAttr(elemType));
rowInit = linalg::MapOp::create(
rewriter, loc, ValueRange {}, rowInit,
[&](OpBuilder& builder, Location nestedLoc, ValueRange) {
Value zero = arith::ConstantOp::create(builder, nestedLoc, zeroAttr);
linalg::YieldOp::create(builder, nestedLoc, zero);
}).getResult().front();
}
auto rowLoop = buildNormalizedScfFor(
rewriter,
loc,
c0,
lanePatches,
c1,
ValueRange {args.inputs[1]},
ValueRange {rowInit},
[&](OpBuilder&, Location nestedLoc, Value copyIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
Value patchIndex = arith::AddIOp::create(rewriter, nestedLoc, laneStart, copyIndex);
Value batchIndex =
affineAddFloorDivConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp);
Value batchIndex = state.batchSize == 1
? c0
: affineAddFloorDivConst(
rewriter, nestedLoc, patchIndex, plan.chunkStart,
plan.numPatchesPerBatch, anchorOp);
Value batchPatchIndex =
affineAddModConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp);
Value outHeightIndex = affineFloorDivConst(rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchorOp);
@@ -2249,7 +2265,8 @@ static Value createIm2colRows(const ConvLoweringState& state,
});
if (failed(rowLoop))
return failure();
publishGraphBatchPhysicalFragment(rewriter, loc, rowLoop->results.front(), args.outputs.front(), args.lane);
Value row = rowLoop->results.front();
publishGraphBatchPhysicalFragment(rewriter, loc, row, args.outputs.front(), args.lane);
return success();
});