/*
 * SPDX-License-Identifier: Apache-2.0
 */

#include "custom_optimizer_passes.h"

#include <memory>
#include <mutex>
#include <string>

#include "onnxoptimizer/optimize.h"
#include "passes/eliminate_nop_dropout.h"
#include "passes/eliminate_reshape_around_elementwise.h"
#include "passes/fuse_add_bias_into_conv.h"
#include "passes/fuse_bn_into_conv.h"
#include "passes/fuse_consecutive_mul.h"
#include "passes/fuse_consecutive_unsqueezes.h"
#include "passes/fuse_matmul_add_bias_into_gemm_batched.h"
#include "passes/fuse_mul_into_conv.h"
#include "passes/fuse_pad_into_pool.h"
#include "passes/fuse_preceding_mul_into_conv.h"

namespace onnxsim {

namespace {

// Register onnxsim's pass T into onnxoptimizer's global registry, replacing any
// existing pass of the same name. onnxoptimizer's own registerPass<T> always
// appends to pass_names, so reusing it to overwrite a name would list -- and
// therefore run -- that pass twice. We touch the registry's public members
// directly instead: overwrite the map entry, and add the name only when it is
// new. This is what lets onnxsim ship its own version of an onnxoptimizer pass
// (e.g. a bug-fixed one) and have it win over the built-in of the same name.
template <typename T>
void RegisterOrReplace(ONNX_NAMESPACE::optimization::GlobalPassRegistry& reg) {
  auto pass = std::make_shared<T>();
  const std::string name = pass->getPassName();
  if (reg.passes.find(name) == reg.passes.end()) {
    reg.pass_names.emplace_back(name);
  }
  reg.passes[name] = std::move(pass);
}

}  // namespace

void RegisterCustomOptimizerPasses() {
  static std::once_flag flag;
  std::call_once(flag, [] {
    namespace p = ONNX_NAMESPACE::optimization::onnxsim_passes;
    // Inject onnxsim's passes into onnxoptimizer's existing global registry,
    // overwriting any built-in of the same name. call_once keeps this a
    // one-time write, matching the registry's static, read-only-after-init use.
    auto& registry = ONNX_NAMESPACE::optimization::Optimizer::passes;

    // onnxsim-only rewrites (no built-in of the same name today).
    RegisterOrReplace<p::EliminateReshapeAroundElementwise>(registry);
    RegisterOrReplace<p::FuseConsecutiveMul>(registry);
    RegisterOrReplace<p::FuseMatMulAddBiasIntoGemmBatched>(registry);
    RegisterOrReplace<p::FuseMulIntoConv>(registry);
    RegisterOrReplace<p::FusePrecedingMulIntoConv>(registry);

    // onnxsim's patched versions of built-in onnxoptimizer passes. These
    // overwrite the registry entries the submodule registers under the same
    // names, so onnxsim's fixes/extensions (ConvTranspose fusions, no-op
    // opset-12 Dropout, zero-padding MaxPool, ...) apply while the fork itself
    // tracks upstream onnxoptimizer.
    RegisterOrReplace<p::EliminateNopDropout>(registry);
    RegisterOrReplace<p::FuseAddBiasIntoConv>(registry);
    RegisterOrReplace<p::FuseBNIntoConv>(registry);
    RegisterOrReplace<p::FuseConsecutiveUnsqueezes>(registry);
    RegisterOrReplace<p::FusePadIntoPool>(registry);
  });
}

}  // namespace onnxsim