forked from ccf-ai-infra/GPUCodeForces
103 lines
2.4 KiB
Python
103 lines
2.4 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
cuda_source = """
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
|
|
__device__ void complex_exp(float re, float im, float& out_re, float& out_im) {
|
|
float exp_re = expf(re);
|
|
out_re = exp_re * cosf(im);
|
|
out_im = exp_re * sinf(im);
|
|
}
|
|
|
|
__device__ void complex_log(float re, float im, float& out_re, float& out_im) {
|
|
out_re = 0.5f * logf(re * re + im * im);
|
|
out_im = atan2f(im, re);
|
|
}
|
|
|
|
__device__ void complex_pow(float z_re, float z_im, float p_re, float p_im, float& out_re, float& out_im) {
|
|
float log_re, log_im;
|
|
complex_log(z_re, z_im, log_re, log_im);
|
|
|
|
float prod_re = p_re * log_re - p_im * log_im;
|
|
float prod_im = p_re * log_im + p_im * log_re;
|
|
|
|
complex_exp(prod_re, prod_im, out_re, out_im);
|
|
}
|
|
|
|
__global__ void complex_ops_kernel(
|
|
const float* __restrict__ z_f,
|
|
float* __restrict__ output,
|
|
float p_re,
|
|
float p_im,
|
|
int batch_size) {
|
|
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (idx < batch_size) {
|
|
float z_re = z_f[idx * 2];
|
|
float z_im = z_f[idx * 2 + 1];
|
|
|
|
float y_re, y_im;
|
|
complex_exp(z_re, z_im, y_re, y_im);
|
|
|
|
float w_re, w_im;
|
|
complex_log(y_re, y_im, w_re, w_im);
|
|
|
|
float out_re, out_im;
|
|
complex_pow(w_re, w_im, p_re, p_im, out_re, out_im);
|
|
|
|
output[idx * 2] = out_re;
|
|
output[idx * 2 + 1] = out_im;
|
|
}
|
|
}
|
|
|
|
torch::Tensor complex_ops_cuda(
|
|
torch::Tensor z_f,
|
|
float p_re,
|
|
float p_im) {
|
|
|
|
int batch_size = z_f.size(0);
|
|
auto output = torch::empty_like(z_f);
|
|
|
|
int threads = 256;
|
|
int blocks = (batch_size + threads - 1) / threads;
|
|
|
|
complex_ops_kernel<<<blocks, threads>>>(
|
|
z_f.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
p_re,
|
|
p_im,
|
|
batch_size
|
|
);
|
|
|
|
return output;
|
|
}
|
|
"""
|
|
|
|
cpp_source = """
|
|
torch::Tensor complex_ops_cuda(
|
|
torch::Tensor z_f,
|
|
float p_re,
|
|
float p_im);
|
|
"""
|
|
|
|
cuda_module = load_inline(
|
|
name="complex_ops_module",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["complex_ops_cuda"],
|
|
verbose=True
|
|
)
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
def __init__(self, p_re, p_im):
|
|
super(ModelNew, self).__init__()
|
|
self.p_re = p_re
|
|
self.p_im = p_im
|
|
|
|
def forward(self, z_f):
|
|
return cuda_module.complex_ops_cuda(z_f, self.p_re, self.p_im) |