#include <torch/extension.h>

#include "cuda_utils.hpp"
#include "engine/pybind.hpp"
#include "graph_timer/graph_timer.hpp"

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    // Cuda utils
    pybind11::module cuda_utils = m.def_submodule("cuda_utils", "DeepInfer cuda utils");
    deepinfer::cuda::utils::pymodule_init(cuda_utils);
    // Engine
    pybind11::module engine = m.def_submodule("engine");
    deepinfer::engine::pymodule_init(engine);
    // CUDA graph 内探针计时
    pybind11::module graph_timer = m.def_submodule("graph_timer", "CUDA graph timer");
    deepinfer::graph_timer::pymodule_init(graph_timer);
}