#!/usr/bin/env python3
"""Prebuild the CUDA helper used verbatim by upstream tests/torch_ref.py.

This runs in the CUDA build container without a GPU. Keeping the extension
name, sources, generated function, and compiler flags identical lets PyTorch
reuse the cache when the official reference module is imported on the host.
"""

from torch.utils.cpp_extension import load_inline


CUDA_SOURCE = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void sigmoid_tanh_fp32_kernel(const float* __restrict__ input,
                                         float* __restrict__ output, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        float xh = input[idx] * 0.5f;
        float th;
        asm("tanh.approx.f32 %0, %1;" : "=f"(th) : "f"(xh));
        output[idx] = th * 0.5f + 0.5f;
    }
}

torch::Tensor sigmoid_tanh_fp32(torch::Tensor input) {
    auto output = torch::empty_like(input);
    int n = input.numel();
    sigmoid_tanh_fp32_kernel<<<(n + 255) / 256, 256>>>(
        input.data_ptr<float>(), output.data_ptr<float>(), n);
    return output;
}
"""


def main() -> None:
    module = load_inline(
        name="sigmoid_ext",
        cpp_sources="torch::Tensor sigmoid_tanh_fp32(torch::Tensor input);",
        cuda_sources=CUDA_SOURCE,
        functions=["sigmoid_tanh_fp32"],
        extra_cuda_cflags=["-O2"],
        verbose=True,
    )
    print(module.__file__)


if __name__ == "__main__":
    main()