#include "mlir/Transforms/Passes.h"
#include "mlir/Conversion/Passes.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
#include "passes.h"
#include "triton/Analysis/Allocation.h"
#include "triton/Analysis/Membar.h"
#include "triton/Conversion/TritonGPUToLLVM/Passes.h"
#include "triton/Conversion/TritonToTritonGPU/Passes.h"
#include "triton/Dialect/Gluon/Transforms/Passes.h"
#include "triton/Dialect/Triton/Transforms/Passes.h"
#include "triton/Dialect/TritonGPU/Transforms/Passes.h"
#include "triton/Dialect/TritonInstrument/Transforms/Passes.h"
#include "triton/Target/LLVMIR/Passes.h"
#include "triton/Tools/PluginUtils.h"
#include <memory>
#include <nanobind/nanobind.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/vector.h>
#include <string>

namespace py = nanobind;

namespace {

void init_triton_analysis(py::module_ &m) {
  py::class_<mlir::ModuleAllocation>(m, "allocation")
      .def(py::init<mlir::ModuleOp>());
  py::class_<mlir::ModuleMembarAnalysis>(m, "membar")
      .def(py::init<mlir::ModuleAllocation &>())
      .def("run", &mlir::ModuleMembarAnalysis::run);
}

void init_triton_passes_common(py::module_ &m) {
  using namespace mlir;
  ADD_PASS_WRAPPER_0("add_sccp", createSCCPPass);
  ADD_PASS_WRAPPER_0("add_symbol_dce", createSymbolDCEPass);
  ADD_PASS_WRAPPER_0("add_inliner", createInlinerPass);
  ADD_PASS_WRAPPER_0("add_canonicalizer", createCanonicalizerPass);
  ADD_PASS_WRAPPER_0("add_cse", createCSEPass);
  ADD_PASS_WRAPPER_0("add_licm", createLoopInvariantCodeMotionPass);
  ADD_PASS_WRAPPER_0("print_ir", createPrintIRPass);
}

void init_triton_passes_ttir(py::module_ &m) {
  using namespace mlir::triton;
  ADD_PASS_WRAPPER_0("add_combine", createTritonCombineOps);
  ADD_PASS_WRAPPER_0("add_reorder_broadcast", createTritonReorderBroadcast);
  ADD_PASS_WRAPPER_0("add_rewrite_tensor_descriptor_to_pointer",
                     createTritonRewriteTensorDescriptorToPointer);
  ADD_PASS_WRAPPER_0("add_loop_unroll", createTritonLoopUnroll);
  ADD_PASS_WRAPPER_0("add_triton_licm", createTritonLoopInvariantCodeMotion);
  ADD_PASS_WRAPPER_0("add_loop_aware_cse", createTritonLoopAwareCSE);
  ADD_PASS_OPTION_WRAPPER_4("add_convert_to_ttgpuir",
                            createConvertTritonToTritonGPU, const std::string &,
                            int, int, int);
}

void init_triton_passes_ttgpuir(py::module_ &m) {
  using namespace mlir;
  using namespace mlir::triton::gpu;
  using namespace mlir::triton::instrument;
  ADD_PASS_WRAPPER_0("add_coalesce", createTritonGPUCoalesce);
  ADD_PASS_WRAPPER_0("add_optimize_thread_locality",
                     createTritonGPUOptimizeThreadLocality);
  ADD_PASS_OPTION_WRAPPER_1("add_hoist_tmem_alloc",
                            createTritonGPUHoistTMEMAlloc, bool);
  ADD_PASS_OPTION_WRAPPER_1("add_assign_latencies",
                            createTritonGPUAssignLatencies, int);
  ADD_PASS_WRAPPER_0("add_schedule_loops", createTritonGPUScheduleLoops);
  ADD_PASS_OPTION_WRAPPER_2("add_pipeline", createTritonGPUPipeline, int, bool);
  ADD_PASS_OPTION_WRAPPER_1("add_warp_specialize",
                            createTritonGPUAutomaticWarpSpecialization, int);
  ADD_PASS_WRAPPER_0("add_prefetch", createTritonGPUPrefetch);
  ADD_PASS_WRAPPER_0("add_accelerate_matmul", createTritonGPUAccelerateMatmul);
  ADD_PASS_WRAPPER_0("add_reorder_instructions",
                     createTritonGPUReorderInstructions);
  ADD_PASS_OPTION_WRAPPER_1("add_f32_dot_tc", createTritonGPUF32DotTC, bool);
  ADD_PASS_OPTION_WRAPPER_1("add_optimize_dot_operands",
                            createTritonGPUOptimizeDotOperands, bool);
  ADD_PASS_WRAPPER_0("add_remove_layout_conversions",
                     createTritonGPURemoveLayoutConversions);
  ADD_PASS_OPTION_WRAPPER_1("add_remove_layout_conversions",
                            createTritonGPURemoveLayoutConversions, bool);
  ADD_PASS_WRAPPER_0("add_reduce_data_duplication",
                     createTritonGPUReduceDataDuplication);
  ADD_PASS_WRAPPER_0("add_allocate_warp_groups",
                     createTritonGPUAllocateWarpGroups);
  ADD_PASS_OPTION_WRAPPER_1("add_allocate_warp_groups",
                            createTritonGPUAllocateWarpGroups, bool);
  ADD_PASS_WRAPPER_0("add_allocate_shared_memory", createAllocateSharedMemory);
  ADD_PASS_WRAPPER_0("add_allocate_global_scratch_memory",
                     createTritonGPUGlobalScratchAllocationPass);
  ADD_PASS_WRAPPER_0("add_combine_tensor_select_and_if",
                     createTritonGPUCombineTensorSelectAndIf);
  ADD_PASS_WRAPPER_0("add_optimize_accumulator_init",
                     createTritonGPUOptimizeAccumulatorInit);
  ADD_PASS_WRAPPER_0("add_fuse_nested_loops", createTritonGPUFuseNestedLoops);
  ADD_PASS_WRAPPER_0("add_coalesce_async_copy",
                     createTritonGPUCoalesceAsyncCopy);
  ADD_PASS_WRAPPER_0("add_global_sanitizer",
                     createTritonInstrumentGlobalSanitizer);
  m.def("add_prepare_consan_captures",
        [](PassManager &pm, const std::string &target) {
          TritonInstrumentPrepareConSanCapturesOptions options;
          options.target = target;
          pm.addPass(createTritonInstrumentPrepareConSanCaptures(options));
        });
  ADD_PASS_WRAPPER_0("add_concurrency_sanitizer",
                     createTritonInstrumentConcurrencySanitizer);
  ADD_PASS_OPTION_WRAPPER_1("add_fp_sanitizer",
                            createTritonInstrumentFpSanitizer, bool);
  ADD_PASS_WRAPPER_0("add_optimize_partition_warps",
                     createTritonGPUOptimizePartitionWarps);
  m.def("add_canonicalize_llvm_ir", [](mlir::PassManager &pm) {
    pm.addNestedPass<mlir::LLVM::LLVMFuncOp>(createCanonicalizeLLVMIR());
  });
}

void init_plugin_passes(py::module_ &m) {
  auto m_ptr = std::make_shared<py::module_>(m);
  m.def(
      "extend_with",
      [m_ptr](const std::string &path) {
        // Load the plugin library.
        auto pluginOrErr = mlir::triton::plugin::TritonPlugin::load(path);
        if (!pluginOrErr) {
          std::string errMsg = llvm::toString(pluginOrErr.takeError());
          throw std::runtime_error(errMsg);
        }
        auto plugin = std::move(*pluginOrErr);

        // Extend this submodule with the passes defined in the plugin.
        py::gil_scoped_acquire acquire;
        for (const auto &pass : plugin.listPasses()) {
          std::string wrapped = std::string("add_") + pass.name;
          m_ptr->def(
              wrapped.c_str(),
              [pass](mlir::PassManager &pm, std::vector<std::string> args) {
                pass.addPass(&pm, args);
              },
              py::arg("pm"), py::arg("args") = std::vector<std::string>());
        }
      },
      "Given a path to a Triton extension, load it and create `add_*` "
      "functions for each pass.");
}

void init_triton_passes_convert(py::module_ &m) {
  using namespace mlir;
  ADD_PASS_WRAPPER_0("add_scf_to_cf", createSCFToControlFlowPass);
  ADD_PASS_WRAPPER_0("add_cf_to_llvmir", createConvertControlFlowToLLVMPass);
  ADD_PASS_WRAPPER_0("add_index_to_llvmir", createConvertIndexToLLVMPass);
  ADD_PASS_WRAPPER_0("add_arith_to_llvmir", createArithToLLVMConversionPass);
  ADD_PASS_WRAPPER_0("add_nvvm_to_llvm", createConvertNVVMToLLVMPass);
  ADD_PASS_WRAPPER_0("add_reconcile_unrealized_casts",
                     createReconcileUnrealizedCastsPass);
}

void init_triton_passes_llvmir(py::module_ &m) {
  using namespace mlir;
  ADD_PASS_WRAPPER_0("add_di_scope", mlir::createLLVMDIScope);
  ADD_PASS_WRAPPER_0("add_di_local_variable", mlir::createLLVMDILocalVariable);
}

void init_gluon_passes(py::module_ &m) {
  using namespace mlir;
  namespace gluon = mlir::triton::gluon;
  ADD_PASS_WRAPPER_0("add_resolve_auto_encodings",
                     gluon::createGluonResolveAutoEncodingsPass);
  ADD_PASS_WRAPPER_0("add_canonicalizer", gluon::createGluonCanonicalize);
  ADD_PASS_WRAPPER_0("add_inliner", gluon::createGluonInline);
  ADD_PASS_WRAPPER_0("add_infer_coalesced_encodings",
                     gluon::createGluonInferCoalescedEncodingsPass);
}

} // namespace

void init_triton_passes(py::module_ &m) {
  auto analysis_m = m.def_submodule("analysis");
  init_triton_analysis(analysis_m);
  auto common_m = m.def_submodule("common");
  init_triton_passes_common(common_m);
  auto convert_m = m.def_submodule("convert");
  init_triton_passes_convert(convert_m);
  auto ttir_m = m.def_submodule("ttir");
  init_triton_passes_ttir(ttir_m);
  auto ttgpuir_m = m.def_submodule("ttgpuir");
  init_triton_passes_ttgpuir(ttgpuir_m);
  auto llvmir_m = m.def_submodule("llvmir");
  init_triton_passes_llvmir(llvmir_m);
  auto gluon_m = m.def_submodule("gluon");
  init_gluon_passes(gluon_m);
  auto plugin_m = m.def_submodule("plugin");
  init_plugin_passes(plugin_m);
}
