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

Unexpected invariant now it's clear (batched in the first tensor rank)
Validate Operations / validate-operations (push) Has been cancelled
Details

Browse Source
This commit is contained in:
ilgeco
2026-07-13 12:05:59 +02:00
parent fed6d343e5
commit 61e3ea9996
29 changed files with 2791 additions and 707 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp
+5 -14
View File
@@ -139,9 +139,7 @@ static RankedTensorType getReducedSliceType(RankedTensorType inputType, ArrayRef
}
static RankedTensorType getLanePackedKeepdimsType(int64_t laneCount, RankedTensorType leafType) {
SmallVector<int64_t> shape(leafType.getShape().begin(), leafType.getShape().end());
shape.front() = laneCount;
return RankedTensorType::get(shape, leafType.getElementType(), leafType.getEncoding());
return spatial::getGraphBatchPhysicalResultType(laneCount, leafType);
}
static SmallVector<int64_t> getKeptAxes(ArrayRef<bool> reducedAxes) {
@@ -191,12 +189,9 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
SmallVector<OpFoldResult> sliceOffsets;
SmallVector<OpFoldResult> sliceSizes;
SmallVector<OpFoldResult> insertOffsets;
SmallVector<OpFoldResult> insertSizes(inputType.getRank(), rewriter.getIndexAttr(1));
SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, inputType.getRank());
sliceOffsets.reserve(inputType.getRank());
sliceSizes.reserve(inputType.getRank());
insertOffsets.reserve(inputType.getRank());
auto batchOp =
createSpatComputeBatch(rewriter,
@@ -209,7 +204,6 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
size_t keptAxisIndex = 0;
sliceOffsets.clear();
sliceSizes.clear();
insertOffsets.clear();
for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
if (isReduced) {
sliceOffsets.push_back(rewriter.getIndexAttr(0));
@@ -224,14 +218,10 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
sliceSizes.push_back(rewriter.getIndexAttr(1));
}
insertOffsets.push_back(args.lane);
insertOffsets.append(inputType.getRank() - 1, rewriter.getIndexAttr(0));
Value slice = tensor::ExtractSliceOp::create(
rewriter, loc, sliceType, args.inputs.front(), sliceOffsets, sliceSizes, unitStrides);
Value reduced = spatial::SpatVAvgOp::create(rewriter, loc, leafType, slice).getResult();
createParallelInsertSliceIntoBatchOutput(
rewriter, loc, reduced, args.outputs.front(), insertOffsets, insertSizes, unitStrides);
publishGraphBatchPhysicalFragment(rewriter, loc, reduced, args.outputs.front(), args.lane);
});
if (failed(batchOp))
return failure();
@@ -276,10 +266,11 @@ static Value buildKeepdimsFromLanePackedBatch(Value batchValue,
if (!pendingLeadingReducedAxes.empty())
expandCompactToKeepdims.back().append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end());
if (batchType.getNumElements() != batchType.getDimSize(0))
return {};
auto reshapeCompute =
createSpatCompute<1>(rewriter, loc, TypeRange {keepdimsType}, {}, ValueRange {batchValue}, [&](Value input) {
auto flatType =
RankedTensorType::get({batchType.getDimSize(0)}, batchType.getElementType(), batchType.getEncoding());
auto flatType = RankedTensorType::get({batchType.getNumElements()}, batchType.getElementType(), batchType.getEncoding());
Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, input, collapseToFlat);
Value compact = flat;
if (compactKeptType != flatType)