huge refactor for high RewritePatterns usage and less ad-hoc cpp code
Validate Operations / validate-operations (push) Has been cancelled
Validate Operations / validate-operations (push) Has been cancelled
remove Spatial many ops in favor of tensor ops like in pim
This commit is contained in:
@@ -5,6 +5,8 @@
|
||||
|
||||
#include "ShapeTilingUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
@@ -30,10 +32,29 @@ SmallVector<Value> sliceTensor(
|
||||
|
||||
for (int64_t i = 0; i < numSlices; i++) {
|
||||
offsets[axis] = rewriter.getIndexAttr(i * sliceSize);
|
||||
if (i == numSlices - 1 && lastSliceSize != 0)
|
||||
int64_t currentSliceSize = sliceSize;
|
||||
if (i == numSlices - 1 && lastSliceSize != 0) {
|
||||
currentSliceSize = lastSliceSize;
|
||||
sizes[axis] = rewriter.getIndexAttr(lastSliceSize);
|
||||
}
|
||||
|
||||
Value slice = tensor::ExtractSliceOp::create(rewriter, loc, tensorToSlice, offsets, sizes, strides);
|
||||
SmallVector<int64_t> sliceShape(shape.begin(), shape.end());
|
||||
sliceShape[axis] = currentSliceSize;
|
||||
auto sliceType =
|
||||
RankedTensorType::get(sliceShape, cast<RankedTensorType>(tensorToSlice.getType()).getElementType());
|
||||
|
||||
Value slice;
|
||||
if (isHostFoldableValue(tensorToSlice)) {
|
||||
slice = tensor::ExtractSliceOp::create(rewriter, loc, tensorToSlice, offsets, sizes, strides);
|
||||
}
|
||||
else {
|
||||
auto sliceCompute =
|
||||
createSpatCompute<1>(rewriter, loc, TypeRange {sliceType}, {}, ValueRange {tensorToSlice}, [&](Value input) {
|
||||
Value computedSlice = tensor::ExtractSliceOp::create(rewriter, loc, input, offsets, sizes, strides);
|
||||
spatial::SpatYieldOp::create(rewriter, loc, computedSlice);
|
||||
});
|
||||
slice = sliceCompute.getResult(0);
|
||||
}
|
||||
slices.push_back(slice);
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user