#pragma once

#include "mlir/Pass/Pass.h"

#include "src/Accelerators/PIM/Pass/PimPasses.hpp"
#include "src/Compiler/CompilerOptions.hpp"

namespace onnx_mlir {

namespace pim {

struct PimBufferizationPass : PassWrapper<PimBufferizationPass, OperationPass<ModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimBufferizationPass)
  StringRef getArgument() const override { return "bufferize-pim"; }
  StringRef getDescription() const override { return "Bufferize PIM and Spatial ops."; }

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

  void runOnOperation() final;

private:
  void annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) const;
};

} // namespace pim

std::unique_ptr<Pass> createBufferizePimPass() { return std::make_unique<pim::PimBufferizationPass>(); }

} // namespace onnx_mlir