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

Merge with fast resnet

Browse Source
This commit is contained in:
ilgeco
2026-07-20 18:01:58 +02:00
parent 5c94e00f43 ab54243fda
commit 961fd613bd
85 changed files with 4475 additions and 4455 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp
+11 -5
View File
@@ -56,12 +56,14 @@ collectFragmentAssemblyCopiesFromBlueprint(spatial::SpatBlueprintOp blueprint,
return blueprint.emitOpError("fragment assembly lowering requires static ranked tensor results");
std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
std::optional<ArrayRef<int64_t>> sourceSlotsAttr = blueprint.getFragmentSourceSlots();
std::optional<ArrayRef<int64_t>> fragmentStridesAttr = blueprint.getFragmentStrides();
if (!operandIndicesAttr || !fragmentStridesAttr)
if (!operandIndicesAttr || !sourceSlotsAttr || !fragmentStridesAttr)
return blueprint.emitOpError(
"fragment assembly lowering requires explicit operand indices and unit strides");
"fragment assembly lowering requires explicit operand indices, source slots, and unit strides");
ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
ArrayRef<int64_t> sourceSlots = *sourceSlotsAttr;
std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
if (!sourceOffsetsAttr)
return blueprint.emitOpError("fragment assembly lowering requires explicit source offsets");
@@ -115,7 +117,11 @@ collectFragmentAssemblyCopiesFromBlueprint(spatial::SpatBlueprintOp blueprint,
copy.sourceType = sourceType;
copy.hostTargetIndex = hostTargetIndex;
copy.lane = lane;
copy.sourceByteOffset = (sourceOffsets[fragmentIndex] + relativeSourceOffset) * static_cast<int64_t>(elementSize);
copy.sourceByteOffset =
(getFragmentAssemblySourceElementOffset(
sourceType, sourceSlots[fragmentIndex], sourceOffsets[fragmentIndex])
+ relativeSourceOffset)
* static_cast<int64_t>(elementSize);
copy.hostByteOffset = hostElementOffset * static_cast<int64_t>(elementSize);
copy.byteSize = chunkElements * static_cast<int64_t>(elementSize);
copies.push_back(copy);
@@ -183,8 +189,8 @@ collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedRe
if (operandIndices[fragmentIndex] != static_cast<int64_t>(use.getOperandNumber()))
continue;
int64_t sourceElementOffset =
sourceSlots[fragmentIndex] * payloadElementCount + sourceOffsets[fragmentIndex];
int64_t sourceElementOffset = getFragmentAssemblySourceElementOffset(
packedResultType, sourceSlots[fragmentIndex], sourceOffsets[fragmentIndex]);
int64_t lane = sourceElementOffset / payloadElementCount;
if (lane < 0 || lane >= static_cast<int64_t>(laneCount))
return failure();
src/PIM/Conversion/SpatialToPim/Common.cpp
+8
View File
@@ -193,6 +193,14 @@ forEachContiguousDestinationChunk(ArrayRef<int64_t> destShape,
return visit(visit, 0);
}
int64_t getFragmentAssemblySourceElementOffset(RankedTensorType sourceType,
int64_t sourceSlot,
int64_t sourceOffset) {
assert(sourceType.getRank() > 0 && sourceType.hasStaticShape()
&& "fragment assembly source must have a static leading slot dimension");
return sourceSlot * (sourceType.getNumElements() / sourceType.getDimSize(0)) + sourceOffset;
}
static mlir::Value
createSteppedOffset(OpBuilder& builder, Location loc, mlir::Value start, mlir::Value index,
int64_t stepBytes, Operation *constantAnchor) {
src/PIM/Conversion/SpatialToPim/Common.hpp
+4
View File
@@ -65,6 +65,10 @@ forEachContiguousDestinationChunk(llvm::ArrayRef<int64_t> destShape,
llvm::function_ref<mlir::LogicalResult(llvm::ArrayRef<int64_t>, int64_t, int64_t)>
callback);
int64_t getFragmentAssemblySourceElementOffset(mlir::RankedTensorType sourceType,
int64_t sourceSlot,
int64_t sourceOffset);
struct FragmentAssemblyCopy {
mlir::Value source;
mlir::RankedTensorType sourceType;
src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp
+8 -3
View File
@@ -44,13 +44,15 @@ static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,
std::optional<StringRef> modeAttr = blueprint.getMode();
std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
std::optional<ArrayRef<int64_t>> sourceSlotsAttr = blueprint.getFragmentSourceSlots();
std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
std::optional<ArrayRef<int64_t>> fragmentStridesAttr = blueprint.getFragmentStrides();
if (!modeAttr || *modeAttr != "fragment_assembly" || !operandIndicesAttr || !sourceOffsetsAttr
|| !fragmentStridesAttr)
if (!modeAttr || *modeAttr != "fragment_assembly" || !operandIndicesAttr || !sourceSlotsAttr
|| !sourceOffsetsAttr || !fragmentStridesAttr)
return blueprint.emitOpError("fragment assembly lowering requires explicit fragment metadata");
ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
ArrayRef<int64_t> sourceSlots = *sourceSlotsAttr;
ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
@@ -102,7 +104,10 @@ static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,
copy.source = source;
copy.sourceType = sourceType;
copy.sourceByteOffset =
(sourceOffsets[fragmentIndex] + relativeSourceOffset) * static_cast<int64_t>(elementSize);
(getFragmentAssemblySourceElementOffset(
sourceType, sourceSlots[fragmentIndex], sourceOffsets[fragmentIndex])
+ relativeSourceOffset)
* static_cast<int64_t>(elementSize);
copy.hostByteOffset = hostElementOffset * static_cast<int64_t>(elementSize);
copy.byteSize = chunkElements * static_cast<int64_t>(elementSize);
copies.push_back(copy);
src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp
+54 -3
View File
@@ -1,7 +1,8 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
#include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -18,6 +19,30 @@ static void copyRaptorDebugAttrs(Operation* source, Operation* target) {
}
}
static Value createDestinationByteOffset(PatternRewriter& rewriter,
tensor::InsertSliceOp insert) {
auto destinationType = cast<RankedTensorType>(insert.getDestType());
SmallVector<int64_t> strides = computeRowMajorStrides(destinationType.getShape());
int64_t elementBytes = getElementTypeSizeInBytes(destinationType.getElementType());
Value total = arith::ConstantIndexOp::create(rewriter, insert.getLoc(), 0);
for (auto [dimension, offset] : llvm::enumerate(insert.getMixedOffsets())) {
int64_t scale = strides[dimension] * elementBytes;
Value component;
if (auto attribute = dyn_cast<Attribute>(offset)) {
component = arith::ConstantIndexOp::create(
rewriter, insert.getLoc(), cast<IntegerAttr>(attribute).getInt() * scale);
} else {
component = cast<Value>(offset);
if (scale != 1)
component = arith::MulIOp::create(
rewriter, insert.getLoc(), component,
arith::ConstantIndexOp::create(rewriter, insert.getLoc(), scale));
}
total = arith::AddIOp::create(rewriter, insert.getLoc(), total, component);
}
return total;
}
struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
using OpRewritePattern::OpRewritePattern;
@@ -40,7 +65,21 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
rewriter.eraseOp(op);
return success();
}
auto outputType = cast<ShapedType>(op.getResult().getType());
auto outputType = cast<RankedTensorType>(op.getResult().getType());
tensor::InsertSliceOp destinationInsert;
if (op->hasOneUse()) {
auto insert = dyn_cast<tensor::InsertSliceOp>(*op->getUsers().begin());
auto destinationType = insert
? dyn_cast<RankedTensorType>(insert.getDestType()) : RankedTensorType();
if (insert && insert.getSource() == op.getOutput()
&& insert.getSourceType() == outputType
&& insert->getBlock() == op->getBlock() && destinationType
&& destinationType.hasStaticShape()
&& isContiguousSubviewWithDynamicOffsets(
destinationType.getShape(), insert.getMixedOffsets(),
insert.getStaticSizes(), insert.getStaticStrides()))
destinationInsert = insert;
}
Value outputBuffer =
tensor::EmptyOp::create(rewriter, op.getLoc(), outputType.getShape(), outputType.getElementType()).getResult();
auto sizeAttr = getTensorSizeInBytesAttr(rewriter, op.getOperation(), op.getResult());
@@ -50,7 +89,19 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, *sizeAttr, op.getSourceCoreId());
copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
Value received = receive.getOutput();
rewriter.replaceOp(op, received);
if (!destinationInsert) {
rewriter.replaceOp(op, received);
return success();
}
rewriter.setInsertionPoint(destinationInsert);
Value targetOffset = createDestinationByteOffset(rewriter, destinationInsert);
Value zero = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
auto copy = pim::PimMemCopyOp::create(
rewriter, op.getLoc(), destinationInsert.getDestType(), targetOffset, zero,
destinationInsert.getDest(), received, *sizeAttr);
rewriter.replaceOp(destinationInsert, copy.getOutput());
rewriter.eraseOp(op);
return success();
}
};
src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp
+7 -3
View File
@@ -608,11 +608,12 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
for (auto [blueprint, operandNumber] : *fragmentAssemblyUses) {
rewriter.setInsertionPointAfterValue(storedValue);
std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
std::optional<ArrayRef<int64_t>> sourceSlotsAttr = blueprint.getFragmentSourceSlots();
std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
std::optional<ArrayRef<int64_t>> stridesAttr = blueprint.getFragmentStrides();
if (!operandIndicesAttr || !sourceOffsetsAttr || !stridesAttr) {
if (!operandIndicesAttr || !sourceSlotsAttr || !sourceOffsetsAttr || !stridesAttr) {
blueprint.emitOpError(
"fragment assembly lowering requires explicit operand, source-offset, and stride metadata");
"fragment assembly lowering requires explicit operand, source-slot, source-offset, and stride metadata");
return ReturnPathLoweringResult::Failure;
}
@@ -626,6 +627,7 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
}
ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
ArrayRef<int64_t> sourceSlots = *sourceSlotsAttr;
ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
@@ -668,7 +670,9 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
elementSize,
producerOp,
"fragment assembly host offset");
auto sourceOffset = getCheckedByteOffset(sourceOffsets[fragmentIndex] + relativeSourceOffset,
int64_t sourceElementOffset = getFragmentAssemblySourceElementOffset(
sourceType, sourceSlots[fragmentIndex], sourceOffsets[fragmentIndex]);
auto sourceOffset = getCheckedByteOffset(sourceElementOffset + relativeSourceOffset,
elementSize,
producerOp,
"fragment assembly source offset");
src/PIM/Conversion/SpatialToPim/SpatialToPim.td
+1 -1
View File
@@ -12,7 +12,7 @@ include "src/Accelerators/PIM/Dialect/Pim/Pim.td"
def spatToPimVMM : Pat<
(SpatVMMOp:$srcOpRes $weight, $vector),
(PimVMMOp $weight, $vector,
(NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
(NativeCodeCall<"tensor::EmptyOp::create($_builder, $_loc, cast<ShapedType>($0.getType()).getShape(), cast<ShapedType>($0.getType()).getElementType())"> $srcOpRes))
>;
def spatToPimVVDMul : Pat<
src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp
+4 -4
View File
@@ -1,4 +1,3 @@
#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
@@ -175,11 +174,11 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
RewritePatternSet coreBodyPatterns(ctx);
populateCoreBodyPatterns(coreBodyPatterns);
populateAffineToStdConversionPatterns(coreBodyPatterns);
FrozenRewritePatternSet frozenCoreBodyPatterns(std::move(coreBodyPatterns));
ConversionTarget coreBodyTarget(*ctx);
coreBodyTarget.addLegalDialect<PimDialect,
coreBodyTarget.addLegalDialect<affine::AffineDialect,
PimDialect,
tensor::TensorDialect,
arith::ArithDialect,
bufferization::BufferizationDialect,
@@ -226,7 +225,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
eraseUnusedTensorPackingOps(funcOp, rewriter);
ConversionTarget communicationTarget(*ctx);
communicationTarget.addLegalDialect<PimDialect,
communicationTarget.addLegalDialect<affine::AffineDialect,
PimDialect,
tensor::TensorDialect,
arith::ArithDialect,
bufferization::BufferizationDialect,