ilgeco
70 lines
3.4 KiB
C++
70 lines
3.4 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 = "nchw_row_strip_fragments";
|
|
|
|
struct RowStripPhysicalValue {
|
|
mlir::Value storage;
|
|
mlir::RankedTensorType logicalType;
|
|
llvm::SmallVector<int64_t, 16> fragmentOffsets;
|
|
llvm::SmallVector<int64_t, 16> fragmentSizes;
|
|
};
|
|
|
|
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(mlir::Value storage,
|
|
mlir::RankedTensorType logicalType,
|
|
mlir::PatternRewriter& rewriter,
|
|
mlir::Location loc);
|
|
|
|
mlir::FailureOr<mlir::Value> applyRowStripRelu(mlir::Value storage,
|
|
mlir::RankedTensorType logicalType,
|
|
mlir::PatternRewriter& rewriter,
|
|
mlir::Location loc);
|
|
|
|
mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(mlir::Value storage,
|
|
mlir::RankedTensorType logicalType,
|
|
mlir::Value bias,
|
|
mlir::PatternRewriter& rewriter,
|
|
mlir::Location loc);
|
|
|
|
} // namespace onnx_mlir
|