forked from ccf-ai-infra/GPUCodeForces
67 lines
1.5 KiB
Python
67 lines
1.5 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>
|
|
|
|
__global__ void logit_sigmoid_shift_kernel(
|
|
const float* __restrict__ input,
|
|
float* __restrict__ output,
|
|
float shift,
|
|
int size
|
|
) {
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (idx < size) {
|
|
float x = input[idx];
|
|
|
|
float y = logf(x / (1.0f - x));
|
|
|
|
|
|
float z = 1.0f / (1.0f + expf(-y));
|
|
|
|
|
|
output[idx] = z + shift;
|
|
}
|
|
}
|
|
|
|
torch::Tensor logit_sigmoid_shift_cuda(torch::Tensor input, float shift) {
|
|
auto output = torch::empty_like(input);
|
|
int size = input.numel();
|
|
|
|
const int block_size = 256;
|
|
int num_blocks = (size + block_size - 1) / block_size;
|
|
|
|
logit_sigmoid_shift_kernel<<<num_blocks, block_size>>>(
|
|
input.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
shift,
|
|
size
|
|
);
|
|
|
|
return output;
|
|
}
|
|
"""
|
|
|
|
cpp_source = """
|
|
torch::Tensor logit_sigmoid_shift_cuda(torch::Tensor input, float shift);
|
|
"""
|
|
|
|
module = load_inline(
|
|
name="logit_sigmoid_shift",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["logit_sigmoid_shift_cuda"],
|
|
verbose=True
|
|
)
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
def __init__(self, shift):
|
|
super(ModelNew, self).__init__()
|
|
self.shift = shift
|
|
self.module = module
|
|
|
|
def forward(self, x):
|
|
return self.module.logit_sigmoid_shift_cuda(x, self.shift) |