#pragma once

#include "mlir/Pass/Pass.h"

#include "src/Dialect/ONNX/ONNXOps.hpp"

namespace onnx_mlir {

using namespace mlir;
extern bool haveSameStaticShape(Value lhs, Value rhs);

namespace spatial {

#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.hpp.inc"

struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ONNXToSpatialPass)
  StringRef getArgument() const override { return "convert-onnx-to-spatial"; }
  StringRef getDescription() const override { return "Lower ONNX ops to Spatial ops."; }

  ONNXToSpatialPass() = default;
  ONNXToSpatialPass(const ONNXToSpatialPass& pass) {}

  void runOnOperation() override;

private:
  void annotateWeightsConstants(func::FuncOp funcOp) const;
};

} // namespace spatial

std::unique_ptr<Pass> createONNXToSpatialPass() { return std::make_unique<spatial::ONNXToSpatialPass>(); }

} // namespace onnx_mlir