NiccoloN
138 lines
5.1 KiB
C++
138 lines
5.1 KiB
C++
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
|
|
#include "mlir/IR/BuiltinAttributes.h"
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
#include "mlir/IR/Location.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/IR/Value.h"
|
|
#include "mlir/Transforms/DialectConversion.h"
|
|
|
|
#include "llvm/ADT/SmallVector.h"
|
|
#include "llvm/ADT/Twine.h"
|
|
#include "llvm/Support/Casting.h"
|
|
|
|
#include <cassert>
|
|
#include <optional>
|
|
#include <utility>
|
|
|
|
#include "Common.hpp"
|
|
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
|
#include "src/Dialect/ONNX/ONNXOps.hpp"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace onnx_mlir {
|
|
|
|
SmallVector<Value> sliceTensor(
|
|
const Value& tensorToSlice, size_t axis, int64_t sliceSize, ConversionPatternRewriter& rewriter, Location loc) {
|
|
ArrayRef<long> shape = getTensorShape(tensorToSlice);
|
|
assert("Invalid axis" && axis < shape.size());
|
|
|
|
SmallVector<OpFoldResult> strides(shape.size(), rewriter.getIndexAttr(1));
|
|
SmallVector<OpFoldResult> offsets(shape.size(), rewriter.getIndexAttr(0));
|
|
SmallVector<OpFoldResult> sizes;
|
|
sizes.reserve(shape.size());
|
|
for (const auto size : shape)
|
|
sizes.push_back(rewriter.getIndexAttr(size));
|
|
sizes[axis] = rewriter.getIndexAttr(sliceSize);
|
|
|
|
long length = shape[axis];
|
|
auto [numSlices, lastSliceSize] = ceilIntegerDivideWithRemainder(length, sliceSize);
|
|
SmallVector<Value> slices;
|
|
slices.reserve(numSlices);
|
|
|
|
for (int64_t i = 0; i < numSlices; i++) {
|
|
offsets[axis] = rewriter.getIndexAttr(i * sliceSize);
|
|
if (i == numSlices - 1 && lastSliceSize != 0)
|
|
sizes[axis] = rewriter.getIndexAttr(lastSliceSize);
|
|
|
|
Value slice = tensor::ExtractSliceOp::create(rewriter, loc, tensorToSlice, offsets, sizes, strides);
|
|
slices.push_back(slice);
|
|
}
|
|
|
|
return slices;
|
|
}
|
|
|
|
SmallVector<Value>
|
|
sliceVector(const Value& vectorToSlice, int64_t sliceSize, ConversionPatternRewriter& rewriter, Location loc) {
|
|
ArrayRef<long> shape = getTensorShape(vectorToSlice);
|
|
assert("Not a vector" && isVectorShape(shape));
|
|
size_t axis = shape[0] != 1 ? 0 : 1;
|
|
return sliceTensor(vectorToSlice, axis, sliceSize, rewriter, loc);
|
|
}
|
|
|
|
DenseMap<CoreId, SmallVector<Value>>
|
|
sliceVectorPerCrossbarPerCore(const Value& vectorToSlice, ConversionPatternRewriter& rewriter, Location loc) {
|
|
SmallVector<Value> slices = sliceVector(vectorToSlice, crossbarSize, rewriter, loc);
|
|
DenseMap<CoreId, SmallVector<Value>> slicesPerCore;
|
|
for (size_t sliceId = 0; sliceId < slices.size(); sliceId++) {
|
|
size_t coreId = sliceId / crossbarCountInCore;
|
|
slicesPerCore[coreId].push_back(slices[sliceId]);
|
|
}
|
|
return slicesPerCore;
|
|
}
|
|
|
|
DenseMap<HSliceId, DenseMap<CoreId, SmallVector<Value>>> tileMatrix(
|
|
Value& matrixToTile, int64_t hSliceSize, int64_t vSliceSize, ConversionPatternRewriter& rewriter, Location& loc) {
|
|
assert("Not a matrix" && isMatrixShape(getTensorShape(matrixToTile)));
|
|
|
|
DenseMap<HSliceId, DenseMap<CoreId, SmallVector<Value>>> tiles;
|
|
|
|
SmallVector<Value> hSlices = sliceTensor(matrixToTile, 1, hSliceSize, rewriter, loc);
|
|
size_t numHSlices = hSlices.size();
|
|
for (size_t hSliceId = 0; hSliceId < numHSlices; hSliceId++) {
|
|
Value hSlice = hSlices[hSliceId];
|
|
SmallVector<Value> vSlices = sliceTensor(hSlice, 0, vSliceSize, rewriter, loc);
|
|
for (size_t vSliceId = 0; vSliceId < vSlices.size(); vSliceId++) {
|
|
size_t coreId = vSliceId / crossbarCountInCore;
|
|
Value vSlice = vSlices[vSliceId];
|
|
tiles[hSliceId][coreId].push_back(vSlice);
|
|
}
|
|
}
|
|
return tiles;
|
|
}
|
|
|
|
tensor::SplatOp
|
|
broadcastToVector(Value scalarToBroadcast, int64_t length, ConversionPatternRewriter& rewriter, Location loc) {
|
|
auto oldType = cast<RankedTensorType>(scalarToBroadcast.getType());
|
|
Type elementType = oldType.getElementType();
|
|
int64_t shape[2] = {1, length};
|
|
Type type = oldType.cloneWith(ArrayRef(shape), elementType);
|
|
|
|
auto zero = arith::ConstantIndexOp::create(rewriter, loc, 0).getResult();
|
|
SmallVector<Value> index(oldType.getRank(), zero);
|
|
auto elementValue = tensor::ExtractOp::create(rewriter, loc, scalarToBroadcast, index).getResult();
|
|
|
|
return tensor::SplatOp::create(rewriter, loc, type, elementValue);
|
|
}
|
|
|
|
Value sumTensors(ArrayRef<Value> tensors, ConversionPatternRewriter& rewriter) {
|
|
if (tensors.size() == 1)
|
|
return tensors[0];
|
|
|
|
SmallVector<Value> tensors1 = {tensors.begin(), tensors.end()};
|
|
SmallVector<Value> tensors2;
|
|
tensors2.reserve(tensors.size() / 2);
|
|
|
|
auto* currTensors = &tensors1;
|
|
auto* nextTensors = &tensors2;
|
|
while (currTensors->size() > 1) {
|
|
for (size_t i = 0; i < currTensors->size() - 1; i += 2) {
|
|
Value a = (*currTensors)[i];
|
|
Value b = (*currTensors)[i + 1];
|
|
rewriter.setInsertionPointAfterValue(b);
|
|
auto addedValue = spatial::SpatVAddOp::create(rewriter, a.getLoc(), a.getType(), a, b);
|
|
nextTensors->push_back(addedValue);
|
|
}
|
|
if (currTensors->size() % 2 == 1)
|
|
nextTensors->push_back(currTensors->back());
|
|
std::swap(currTensors, nextTensors);
|
|
nextTensors->clear();
|
|
}
|
|
assert(currTensors->size() == 1 && "Expected a single input at this point.");
|
|
return (*currTensors)[0];
|
|
}
|
|
|
|
}; // namespace onnx_mlir
|