wuyang
51 lines
1.4 KiB
Python
51 lines
1.4 KiB
Python
#!/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()
|