NiccoloN
120 lines
3.4 KiB
C++
120 lines
3.4 KiB
C++
#include "mlir/IR/ValueRange.h"
|
|
|
|
#include "llvm/ADT/STLExtras.h"
|
|
#include "llvm/ADT/StringRef.h"
|
|
|
|
#include <cassert>
|
|
#include <cstddef>
|
|
|
|
#include "Common.hpp"
|
|
|
|
using namespace llvm;
|
|
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;
|
|
}
|
|
|
|
size_t getShapedTypeSizeInBytes(ShapedType shapedType) {
|
|
return shapedType.getNumElements() * shapedType.getElementTypeBitWidth() / 8;
|
|
}
|
|
|
|
IntegerAttr getTensorSizeInBytesAttr(Builder& builder, mlir::Value value) {
|
|
return builder.getI32IntegerAttr(static_cast<int32_t>(getShapedTypeSizeInBytes(cast<ShapedType>(value.getType()))));
|
|
}
|
|
|
|
Operation* getEarliestUserWithinBlock(mlir::Value value) {
|
|
auto users = value.getUsers();
|
|
|
|
assert(!users.empty());
|
|
|
|
Operation* earliestUser = *users.begin();
|
|
for (auto curUser : users)
|
|
if (curUser->isBeforeInBlock(earliestUser))
|
|
earliestUser = curUser;
|
|
|
|
return earliestUser;
|
|
}
|
|
|
|
SmallVector<mlir::Value> getOpOperandsSortedByUses(Operation* operation) {
|
|
auto operandsAndUses =
|
|
map_to_vector(operation->getOperands(), [](mlir::Value operand) -> std::pair<mlir::Value, size_t> {
|
|
return {operand, std::distance(operand.use_begin(), operand.use_end())};
|
|
});
|
|
sort(operandsAndUses, [](auto a, auto b) { return a.second < b.second; });
|
|
return map_to_vector(operandsAndUses, [](auto operandAndUse) { return operandAndUse.first; });
|
|
}
|
|
|
|
bool hasLaterUserInBlock(mlir::Value value, Operation* operation) {
|
|
for (Operation* user : value.getUsers()) {
|
|
if (user->getBlock() != operation->getBlock())
|
|
return true;
|
|
if (operation->isBeforeInBlock(user))
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
mlir::Value getBestOutputTensorFromOperandsOrAllocate(PatternRewriter& rewriter, Operation* operation) {
|
|
assert("Only support operations with a single result" && operation->getNumResults() == 1);
|
|
mlir::Value result = operation->getResult(0);
|
|
auto resultType = result.getType();
|
|
assert("Only support result ShapedType as result type" && isa<ShapedType>(resultType));
|
|
|
|
SmallVector<mlir::Value> operands = getOpOperandsSortedByUses(operation);
|
|
auto validOperands = make_filter_range(operands, [operation, resultType](mlir::Value operand) {
|
|
return operand.getType() == resultType && !hasLaterUserInBlock(operand, operation);
|
|
});
|
|
auto bestOperand = validOperands.begin();
|
|
|
|
if (bestOperand != validOperands.end())
|
|
return *bestOperand;
|
|
|
|
auto resultShapedType = cast<ShapedType>(resultType);
|
|
rewriter.setInsertionPoint(operation);
|
|
return tensor::EmptyOp::create(
|
|
rewriter, operation->getLoc(), resultShapedType.getShape(), resultShapedType.getElementType());
|
|
}
|
|
|
|
} // namespace onnx_mlir
|