// ix_full_bridge_fused_ar.cpp — Bridge to ixformer C++ functions + fused linear+allreduce
//
// Adds ixformer_linear_allreduce (linear + NCCL allreduce fused in one kernel launch)
//
// Symbols from _ixformer_torch.cpython-310-x86_64-linux-gnu.so (nm -D verified):
//   ixformer_torch_ext::silu_and_mul_forward(at::Tensor&, at::Tensor&)
//   ixformer_torch_ext::rms_norm_forward(at::Tensor&, at::Tensor&, at::Tensor&, double)
//   ixformer_torch_ext::fused_add_rms_norm_forward(at::Tensor&, at::Tensor&, at::Tensor&, double, double)
//   ixformer_torch_ext::ixformer_linear(at::Tensor&, at::Tensor&, optional<Tensor>, optional<Tensor>)
//   ixformer_torch_ext::ixformer_linear_ex(at::Tensor&, at::Tensor&, optional<Tensor>)
//   ixformer_torch_ext::ixformer_linear_allreduce(at::Tensor&, at::Tensor&, optional<Tensor>)  ← NEW

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

namespace ixformer_torch_ext {

void silu_and_mul_forward(at::Tensor& input, at::Tensor& output);

void rms_norm_forward(at::Tensor& input, at::Tensor& weight,
                      at::Tensor& output, double eps);

void fused_add_rms_norm_forward(at::Tensor& input, at::Tensor& residual,
                                at::Tensor& weight, double eps, double alpha);

at::Tensor ixformer_linear(at::Tensor& input, at::Tensor& weight,
                           const c10::optional<at::Tensor>& bias,
                           const c10::optional<at::Tensor>& out);

at::Tensor ixformer_linear_ex(at::Tensor& input, at::Tensor& weight,
                              const c10::optional<at::Tensor>& bias);

// fused linear + allreduce
// Signature from nm -D on _ixformer_torch.cpython-310-x86_64-linux-gnu.so:
//   000000000003c410 T ixformer_torch_ext::ixformer_linear_allreduce(at::Tensor&, at::Tensor&, c10::optional<at::Tensor> const&)
at::Tensor ixformer_linear_allreduce(at::Tensor& input, at::Tensor& weight,
                                     const c10::optional<at::Tensor>& bias);

}  // namespace ixformer_torch_ext


// ============================================================================
// Python wrappers
// ============================================================================

torch::Tensor ix_silu_and_mul(torch::Tensor input) {
  int64_t half_dim = input.size(-1) / 2;
  auto output = input.new_empty({input.size(0), half_dim});
  ixformer_torch_ext::silu_and_mul_forward(input, output);
  return output;
}

void ix_rms_norm(torch::Tensor output, torch::Tensor input,
                 torch::Tensor weight, double eps) {
  ixformer_torch_ext::rms_norm_forward(input, weight, output, eps);
}

void ix_fused_add_rms_norm(torch::Tensor input, torch::Tensor residual,
                           torch::Tensor weight, double eps) {
  ixformer_torch_ext::fused_add_rms_norm_forward(input, residual, weight, eps, 1.0);
}

torch::Tensor ix_linear(torch::Tensor input, torch::Tensor weight,
                         const c10::optional<torch::Tensor>& bias) {
  auto input_2d = input.view({-1, input.size(-1)});
  int64_t m = input_2d.size(0);
  if (m <= 1 && !bias.has_value()) {
    return ixformer_torch_ext::ixformer_linear_ex(input, weight, bias);
  }
  return ixformer_torch_ext::ixformer_linear(input, weight, bias,
                                              c10::optional<at::Tensor>());
}

// fused linear + allreduce — single kernel launch
// Replaces: output = linear(input, weight); output = allreduce(output)
torch::Tensor ix_linear_allreduce(torch::Tensor input, torch::Tensor weight,
                                   const c10::optional<torch::Tensor>& bias) {
  return ixformer_torch_ext::ixformer_linear_allreduce(input, weight, bias);
}


PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("silu_and_mul", &ix_silu_and_mul, "Fused SiLU+mul activation");
  m.def("rms_norm", &ix_rms_norm, "RMSNorm");
  m.def("fused_add_rms_norm", &ix_fused_add_rms_norm, "Fused residual + RMSNorm");
  m.def("linear", &ix_linear, "ixformer GEMM (linear/linear_ex auto-dispatch)");
  m.def("linear_allreduce", &ix_linear_allreduce,
        "Fused linear + allreduce (saves one kernel launch per RowParallel op)");
}