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

centralize logic for materializing contiguous memory into bufferization

Browse Source 
fix codegen symlinks overwrite
remove deprecated pim memcp_hd_batch op
This commit is contained in:
NiccoloN
2026-05-30 15:54:24 +02:00
parent 2d5b03c08f
commit ff36729140
29 changed files with 642 additions and 822 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp
+19 -20
View File
@@ -59,10 +59,7 @@ static Value createHostTargetOffset(IRRewriter& rewriter,
ShapedType destinationType,
IRMapping& mapper) {
int64_t elementBytes = static_cast<int64_t>(getElementTypeSizeInBytes(destinationType.getElementType()));
SmallVector<int64_t> strides(destinationType.getRank(), 1);
ArrayRef<int64_t> shape = destinationType.getShape();
for (int64_t dim = destinationType.getRank() - 2; dim >= 0; --dim)
strides[dim] = strides[dim + 1] * shape[dim + 1];
SmallVector<int64_t> strides = computeRowMajorStrides(destinationType.getShape());
Value totalOffset;
Location loc = insertSlice.getLoc();
@@ -162,14 +159,15 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
BlockArgument newArg = coreBatchOp.getInputArgument(inputIndex);
auto newArgType = cast<ShapedType>(newArg.getType());
auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, newArgType);
auto copied = pim::PimMemCopyHostToDevBatchOp::create(rewriter,
loc,
outputBuffer.getType(),
outputBuffer,
newArg,
rewriter.getI32IntegerAttr(0),
rewriter.getI32IntegerAttr(0),
getTensorSizeInBytesAttr(rewriter, newArg))
Value zeroOffset = getOrCreateIndexConstant(rewriter, coreBatchOp.getOperation(), 0);
auto copied = pim::PimMemCopyHostToDevOp::create(rewriter,
loc,
outputBuffer.getType(),
zeroOffset,
zeroOffset,
outputBuffer,
newArg,
getTensorSizeInBytesAttr(rewriter, newArg))
.getOutput();
mapper.map(*oldArg, copied);
}
@@ -233,14 +231,15 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
}
auto clonedType = cast<ShapedType>(clonedTensor.getType());
auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, clonedType);
auto copied = pim::PimMemCopyHostToDevBatchOp::create(rewriter,
loc,
outputBuffer.getType(),
outputBuffer,
clonedTensor,
rewriter.getI32IntegerAttr(0),
rewriter.getI32IntegerAttr(0),
getTensorSizeInBytesAttr(rewriter, clonedTensor))
Value zeroOffset = getOrCreateIndexConstant(rewriter, coreBatchOp.getOperation(), 0);
auto copied = pim::PimMemCopyHostToDevOp::create(rewriter,
loc,
outputBuffer.getType(),
zeroOffset,
zeroOffset,
outputBuffer,
clonedTensor,
getTensorSizeInBytesAttr(rewriter, clonedTensor))
.getOutput();
mapper.map(toTensorOp.getResult(), copied);
continue;
src/PIM/Conversion/SpatialToPim/Common.cpp
-44
View File
@@ -1,10 +1,8 @@
#include "mlir/IR/ValueRange.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/StringRef.h"
#include <cassert>
#include <cstddef>
#include "Common.hpp"
@@ -13,48 +11,6 @@ using namespace mlir;
namespace onnx_mlir {
size_t getSliceActualOffset(tensor::ExtractSliceOp& sliceOp, ShapedType& inputShape) {
/*
EXAMPLE RUN:
[1, 10, 3, 4] inputShape
[0, 2, 1, 3] offsets
acc = 1
---
ret = 3
acc = 4
---
ret = 3 + 4 * 1 = 7
acc = 12
---
ret = 7 + 12 * 2 = 31
acc = 120
---
ret = 31 + 120 * 0 = 31
acc = 120
*/
size_t returnValue = 0;
auto sliceOffsets = sliceOp.getStaticOffsets();
auto inputDimSizes = inputShape.getShape();
assert(sliceOffsets.size() == inputDimSizes.size());
size_t accumulatedDimensionSize = 1;
// Reverse iterate the two vectors
for (auto it : reverse(zip(sliceOffsets, inputDimSizes))) {
auto curSliceOffset = std::get<0>(it);
auto curInputDimSize = std::get<1>(it);
returnValue += accumulatedDimensionSize * curSliceOffset;
accumulatedDimensionSize *= curInputDimSize;
}
return returnValue;
}
IntegerAttr getTensorSizeInBytesAttr(Builder& builder, mlir::Value value) {
return builder.getI32IntegerAttr(static_cast<int32_t>(getShapedTypeSizeInBytes(cast<ShapedType>(value.getType()))));
}
src/PIM/Conversion/SpatialToPim/Common.hpp
-14
View File
@@ -6,20 +6,6 @@
namespace onnx_mlir {
/**
* \brief Get the offset of the ExtractSliceOp based on its static offsets and
* its static tensor input.
*
* The static offsets represent the starting position of the slice in each
* dimension, while the static tensor input gives its dimension size.
*
* \param sliceOp The ExtractSliceOp for which the actual offset needs to be
* calculated.
* \param inputShape The ShapedType of the ExtractSliceOp's input tensor
* \return The actual offset of the ExtractSliceOp.
*/
size_t getSliceActualOffset(mlir::tensor::ExtractSliceOp& sliceOp, mlir::ShapedType& inputShape);
mlir::IntegerAttr getTensorSizeInBytesAttr(mlir::Builder& builder, mlir::Value value);
template <class T>
src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp
-12
View File
@@ -83,18 +83,6 @@ static Value createZeroedDeviceHVector(IRRewriter& rewriter,
auto zeroValue = memref::GetGlobalOp::create(rewriter, loc, zeroGlobal.getType(), zeroGlobal.getName());
auto zeroIndex = getOrCreateIndexConstant(constantFolder, outputBuffer.getOperation(), 0);
auto sizeAttr = rewriter.getI32IntegerAttr(static_cast<int32_t>(getShapedTypeSizeInBytes(tensorType)));
if (outputBuffer->getParentOfType<PimCoreBatchOp>())
return PimMemCopyHostToDevBatchOp::create(rewriter,
loc,
tensorType,
outputBuffer,
zeroValue,
rewriter.getI32IntegerAttr(0),
rewriter.getI32IntegerAttr(0),
sizeAttr)
.getOutput();
return PimMemCopyHostToDevOp::create(
rewriter, loc, tensorType, zeroIndex, zeroIndex, outputBuffer, zeroValue, sizeAttr)
.getOutput();