NiccoloN
80 lines
3.9 KiB
C++
80 lines
3.9 KiB
C++
#pragma once
|
|
|
|
#include "mlir/IR/BuiltinAttributes.h"
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
|
|
namespace onnx_mlir {
|
|
|
|
inline constexpr llvm::StringLiteral kRowStripIndexMap = "nhwc_row_strip_fragments";
|
|
|
|
struct RowStripPhysicalValue {
|
|
mlir::Value storage;
|
|
mlir::RankedTensorType logicalType;
|
|
mlir::RankedTensorType fragmentType;
|
|
int64_t tilesPerRow;
|
|
};
|
|
|
|
mlir::FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(mlir::Value storage,
|
|
mlir::RankedTensorType logicalType);
|
|
|
|
std::pair<llvm::SmallVector<int64_t>, llvm::SmallVector<int64_t>>
|
|
buildRowStripMetadata(mlir::RankedTensorType type);
|
|
|
|
mlir::RankedTensorType getRowStripFragmentType(mlir::RankedTensorType logicalType);
|
|
|
|
mlir::RankedTensorType getRowStripStorageType(mlir::RankedTensorType logicalType);
|
|
|
|
llvm::SmallVector<mlir::OpFoldResult> buildRowStripFragmentOffsets(mlir::PatternRewriter& rewriter,
|
|
mlir::OpFoldResult row);
|
|
|
|
llvm::SmallVector<mlir::OpFoldResult> buildRowStripFragmentSizes(mlir::PatternRewriter& rewriter,
|
|
mlir::RankedTensorType logicalType);
|
|
|
|
mlir::Value extractRowStripFragment(mlir::Value storage,
|
|
mlir::RankedTensorType logicalType,
|
|
mlir::OpFoldResult row,
|
|
mlir::PatternRewriter& rewriter,
|
|
mlir::Location loc);
|
|
|
|
void insertRowStripFragment(mlir::Value fragment,
|
|
mlir::Value output,
|
|
mlir::RankedTensorType logicalType,
|
|
mlir::OpFoldResult row,
|
|
mlir::PatternRewriter& rewriter,
|
|
mlir::Location loc);
|
|
|
|
mlir::FailureOr<mlir::Value> createPerChannelConstantFragment(mlir::DenseElementsAttr denseAttr,
|
|
mlir::RankedTensorType fragmentType,
|
|
mlir::PatternRewriter& rewriter);
|
|
|
|
mlir::FailureOr<mlir::Value> createRowStripStorageFromRows(mlir::Value rows,
|
|
mlir::RankedTensorType logicalType,
|
|
mlir::PatternRewriter& rewriter,
|
|
mlir::Location loc);
|
|
|
|
mlir::FailureOr<mlir::Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
|
|
mlir::PatternRewriter& rewriter,
|
|
mlir::Location loc);
|
|
|
|
mlir::FailureOr<mlir::Value> applyRowStripRelu(const RowStripPhysicalValue& value,
|
|
mlir::PatternRewriter& rewriter,
|
|
mlir::Location loc);
|
|
|
|
mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
|
|
mlir::Value bias,
|
|
mlir::PatternRewriter& rewriter,
|
|
mlir::Location loc);
|
|
|
|
mlir::FailureOr<mlir::Value> applyRowStripAdd(const RowStripPhysicalValue& lhs,
|
|
const RowStripPhysicalValue& rhs,
|
|
mlir::PatternRewriter& rewriter,
|
|
mlir::Location loc);
|
|
|
|
mlir::FailureOr<mlir::Value> applyRowStripConcat(llvm::ArrayRef<RowStripPhysicalValue> inputs,
|
|
mlir::RankedTensorType outputType,
|
|
mlir::PatternRewriter& rewriter,
|
|
mlir::Location loc);
|
|
|
|
} // namespace onnx_mlir
|