#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