ilgeco
111 lines
4.5 KiB
C++
111 lines
4.5 KiB
C++
#pragma once
|
|
|
|
#include <optional>
|
|
|
|
#include "llvm/ADT/ArrayRef.h"
|
|
#include "llvm/ADT/SmallVector.h"
|
|
#include "llvm/ADT/STLFunctionalExtras.h"
|
|
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
#include "mlir/IR/Builders.h"
|
|
#include "mlir/IR/Value.h"
|
|
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/Support/LogicalResult.h"
|
|
|
|
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
|
|
|
|
namespace onnx_mlir::spatial {
|
|
class SpatBlueprintOp;
|
|
}
|
|
|
|
namespace onnx_mlir {
|
|
|
|
mlir::FailureOr<mlir::IntegerAttr>
|
|
getTensorSizeInBytesAttr(mlir::Builder& builder, mlir::Operation* anchor, mlir::Value value);
|
|
|
|
template <class T>
|
|
size_t rangeLength(const mlir::iterator_range<T> range) {
|
|
return std::distance(range.begin(), range.end());
|
|
}
|
|
|
|
/**
|
|
* Retrieves the earliest operation that uses the given value within the value's
|
|
* block.
|
|
*
|
|
* @param value The value for which to find the earliest user operation.
|
|
* @return The earliest user operation that uses the given value within the
|
|
* current block.
|
|
*/
|
|
mlir::Operation* getEarliestUserWithinBlock(mlir::Value value);
|
|
|
|
mlir::SmallVector<mlir::Value> getOpOperandsSortedByUses(mlir::Operation* operation);
|
|
|
|
mlir::Value getBestOutputTensorFromOperandsOrAllocate(mlir::RewriterBase& rewriter, mlir::Operation* operation);
|
|
|
|
mlir::LogicalResult validateFragmentAssemblyMetadata(onnx_mlir::spatial::SpatBlueprintOp blueprint,
|
|
int64_t resultRank,
|
|
size_t operandCount,
|
|
llvm::ArrayRef<int64_t> operandIndices,
|
|
llvm::ArrayRef<int64_t> sourceOffsets,
|
|
llvm::ArrayRef<int64_t> flatOffsets,
|
|
llvm::ArrayRef<int64_t> flatSizes,
|
|
llvm::ArrayRef<int64_t> flatStrides);
|
|
|
|
mlir::FailureOr<mlir::SmallVector<int64_t, 4>>
|
|
getStaticSliceOffsetsForElementOffset(mlir::Operation* anchor,
|
|
mlir::ShapedType sourceType,
|
|
llvm::ArrayRef<int64_t> fragmentShape,
|
|
int64_t sourceElementOffset,
|
|
llvm::StringRef fieldName);
|
|
|
|
mlir::LogicalResult
|
|
forEachContiguousDestinationChunk(llvm::ArrayRef<int64_t> destShape,
|
|
llvm::ArrayRef<int64_t> baseOffsets,
|
|
llvm::ArrayRef<int64_t> sizes,
|
|
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;
|
|
unsigned hostTargetIndex = 0;
|
|
int64_t lane = 0;
|
|
int64_t sourceByteOffset = 0;
|
|
int64_t hostByteOffset = 0;
|
|
int64_t byteSize = 0;
|
|
};
|
|
|
|
struct FragmentAssemblyCopyRun {
|
|
mlir::Value source;
|
|
mlir::RankedTensorType sourceType;
|
|
unsigned hostTargetIndex = 0;
|
|
int64_t count = 0;
|
|
int64_t sourceStepBytes = 0;
|
|
int64_t hostStepBytes = 0;
|
|
int64_t byteSize = 0;
|
|
mlir::SmallVector<int64_t, 8> sourceStartBytesByLane;
|
|
mlir::SmallVector<int64_t, 8> hostStartBytesByLane;
|
|
};
|
|
|
|
mlir::FailureOr<mlir::SmallVector<FragmentAssemblyCopyRun, 8>>
|
|
groupFragmentAssemblyCopyRuns(llvm::ArrayRef<FragmentAssemblyCopy> copies, uint32_t laneCount = 1);
|
|
|
|
mlir::FailureOr<mlir::Value> emitFragmentAssemblyCopyRuns(mlir::IRRewriter& rewriter,
|
|
mlir::Location loc,
|
|
llvm::ArrayRef<FragmentAssemblyCopyRun> runs,
|
|
mlir::Value hostTarget,
|
|
mlir::Operation* anchor,
|
|
std::optional<mlir::Value> laneArg = std::nullopt,
|
|
mlir::Value baseHostOffset = {});
|
|
|
|
inline mlir::tensor::EmptyOp
|
|
createEmptyTensorFromShaped(mlir::IRRewriter& rewriter, mlir::Location loc, mlir::ShapedType shapedType) {
|
|
return mlir::tensor::EmptyOp::create(rewriter, loc, shapedType.getShape(), shapedType.getElementType());
|
|
}
|
|
|
|
} // namespace onnx_mlir
|