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

finally faster on all pimcomp models
Validate Operations / validate-operations (push) Has been cancelled
Details

Browse Source
This commit is contained in:
NiccoloN
2026-07-30 14:07:01 +02:00
parent 1b4f070bef
commit c12f69133d
4 changed files with 140 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp
+48 -8
View File
@@ -9,6 +9,8 @@
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"
#include <numeric>
using namespace mlir;
namespace onnx_mlir {
@@ -255,10 +257,18 @@ FailureOr<Value> applyRowStripAdd(const RowStripPhysicalValue& lhs,
const RowStripPhysicalValue& rhs,
PatternRewriter& rewriter,
Location loc) {
if (lhs.logicalType != rhs.logicalType || lhs.fragmentType != rhs.fragmentType
|| lhs.storage.getType() != rhs.storage.getType() || lhs.tilesPerRow != rhs.tilesPerRow)
if (lhs.logicalType != rhs.logicalType)
return failure();
auto storageType = cast<RankedTensorType>(lhs.storage.getType());
const int64_t fragmentChannels =
std::gcd(lhs.fragmentType.getDimSize(3), rhs.fragmentType.getDimSize(3));
auto fragmentType = RankedTensorType::get(
{1, 1, lhs.logicalType.getDimSize(3), fragmentChannels},
lhs.logicalType.getElementType(),
lhs.logicalType.getEncoding());
const int64_t tilesPerRow = ceilIntegerDivide(lhs.logicalType.getDimSize(1), fragmentChannels);
auto storageType =
spatial::getGraphBatchPhysicalResultType(lhs.logicalType.getDimSize(2) * tilesPerRow, fragmentType);
const int64_t laneCount = storageType.getDimSize(0);
auto batch = createSpatComputeBatch(
rewriter,
@@ -268,13 +278,43 @@ FailureOr<Value> applyRowStripAdd(const RowStripPhysicalValue& lhs,
{},
ValueRange {lhs.storage, rhs.storage},
[&](detail::SpatComputeBatchBodyArgs args) {
FailureOr<Value> lhsFragment =
extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, lhs.fragmentType);
FailureOr<Value> rhsFragment =
extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[1], args.lane, rhs.fragmentType);
auto extractFragment = [&](Value storage, const RowStripPhysicalValue& input) -> FailureOr<Value> {
Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
MLIRContext* context = rewriter.getContext();
AffineExpr lane = getAffineDimExpr(0, context);
AffineExpr outputTile = lane % tilesPerRow;
AffineExpr channelOffset = outputTile * fragmentChannels;
Value sourceSlot = createOrFoldAffineApply(
rewriter,
loc,
lane.floorDiv(tilesPerRow) * input.tilesPerRow
+ channelOffset.floorDiv(input.fragmentType.getDimSize(3)),
ValueRange {args.lane},
anchor);
FailureOr<Value> source =
extractGraphBatchPhysicalFragment(rewriter, loc, storage, sourceSlot, input.fragmentType);
if (failed(source) || input.fragmentType == fragmentType)
return source;
Value sourceOffset = createOrFoldAffineApply(
rewriter, loc, channelOffset % input.fragmentType.getDimSize(3), ValueRange {args.lane}, anchor);
return tensor::ExtractSliceOp::create(
rewriter,
loc,
fragmentType,
*source,
SmallVector<OpFoldResult> {
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), sourceOffset},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(1),
rewriter.getIndexAttr(lhs.logicalType.getDimSize(3)),
rewriter.getIndexAttr(fragmentChannels)},
getUnitStrides(rewriter, 4)).getResult();
};
FailureOr<Value> lhsFragment = extractFragment(args.inputs[0], lhs);
FailureOr<Value> rhsFragment = extractFragment(args.inputs[1], rhs);
if (failed(lhsFragment) || failed(rhsFragment))
return failure();
Value added = spatial::SpatVAddOp::create(rewriter, loc, lhs.fragmentType, *lhsFragment, *rhsFragment);
Value added = spatial::SpatVAddOp::create(rewriter, loc, fragmentType, *lhsFragment, *rhsFragment);
publishGraphBatchPhysicalFragment(rewriter, loc, added, args.outputs.front(), args.lane);
return success();
});