/*
 * flux_ext.cpp — TORCH_LIBRARY registration for FLUXNATION
 *
 * Registers ssb_prepare and ssb_finish as torch.ops.flux_ext.* custom ops.
 * This eliminates ctypes overhead and enables torch.compile compatibility.
 */

#include <torch/extension.h>
#include <ATen/ATen.h>

// Forward declarations of CUDA implementations (in flux_ext_kernels.cu)
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor>
ssb_prepare_impl(
    at::Tensor x, at::Tensor vec, at::Tensor pe,
    at::Tensor mod_w, at::Tensor mod_w_scale, at::Tensor mod_b,
    at::Tensor l1w, at::Tensor l1w_scale, at::Tensor l1b,
    at::Tensor qns, at::Tensor kns);

at::Tensor ssb_finish_impl(
    at::Tensor x, at::Tensor attn, at::Tensor mlp, at::Tensor gate,
    at::Tensor l2w, at::Tensor l2w_scale, at::Tensor l2b);

void ssb_cleanup_impl();

// =========================================================================
// Op schema definitions
// =========================================================================
TORCH_LIBRARY(flux_ext, m) {
    m.def("ssb_prepare("
          "Tensor x, Tensor vec, Tensor pe, "
          "Tensor mod_w, Tensor mod_w_scale, Tensor mod_b, "
          "Tensor l1w, Tensor l1w_scale, Tensor l1b, "
          "Tensor qns, Tensor kns"
          ") -> (Tensor, Tensor, Tensor, Tensor, Tensor)");

    m.def("ssb_finish("
          "Tensor x, Tensor attn, Tensor mlp, Tensor gate, "
          "Tensor l2w, Tensor l2w_scale, Tensor l2b"
          ") -> Tensor");

    m.def("ssb_cleanup() -> ()");
}

// =========================================================================
// CUDA backend implementations
// =========================================================================
TORCH_LIBRARY_IMPL(flux_ext, CUDA, m) {
    m.impl("ssb_prepare", &ssb_prepare_impl);
    m.impl("ssb_finish",  &ssb_finish_impl);
    m.impl("ssb_cleanup", &ssb_cleanup_impl);
}

// =========================================================================
// Python module init (required for import flux_ext to work)
// =========================================================================
PYBIND11_MODULE(flux_ext, m) {
    m.doc() = "FLUXNATION: Fused CUDA kernels for FLUX.1 SingleStreamBlock";
}