31 lines
1.3 KiB
C++
31 lines
1.3 KiB
C++
#pragma once
|
|
|
|
#include "mlir/IR/BuiltinAttributes.h"
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/IR/Value.h"
|
|
#include "mlir/Support/LogicalResult.h"
|
|
|
|
namespace onnx_mlir {
|
|
|
|
struct BiasAddPlanCandidate {
|
|
mlir::Value data;
|
|
mlir::Value bias;
|
|
};
|
|
|
|
mlir::LogicalResult isSupportedBiasAddShape(mlir::RankedTensorType biasType, mlir::RankedTensorType resultType);
|
|
bool isSupportedBiasAddValue(mlir::Value bias,
|
|
mlir::RankedTensorType resultType,
|
|
mlir::DenseElementsAttr* denseAttr = nullptr);
|
|
mlir::FailureOr<llvm::SmallVector<mlir::Attribute>>
|
|
getBiasChannelValues(mlir::DenseElementsAttr denseAttr, mlir::RankedTensorType resultType);
|
|
mlir::FailureOr<BiasAddPlanCandidate> classifyBiasAddPlanCandidate(mlir::Value lhs,
|
|
mlir::Value rhs,
|
|
mlir::RankedTensorType resultType);
|
|
mlir::FailureOr<mlir::Value> materializeDenseBiasAddTensor(mlir::Value bias,
|
|
mlir::RankedTensorType resultType,
|
|
mlir::RewriterBase& rewriter,
|
|
mlir::Location loc);
|
|
|
|
} // namespace onnx_mlir
|