forked from ccf-ai-infra/GPUCodeForces
288 lines
8.7 KiB
Python
288 lines
8.7 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
N, C, H, W = 32, 1, 64, 64
|
|
|
|
|
|
class GeneratorLossCUDAOp(torch.autograd.Function):
|
|
|
|
def forward(ctx, input, target, beta, reduction_id, op):
|
|
if not input.is_cuda: input = input.cuda()
|
|
if not target.is_cuda: target = target.cuda()
|
|
|
|
input = input.contiguous()
|
|
target = target.contiguous()
|
|
|
|
n = input.numel()
|
|
|
|
output, partial_sums = op.generator_loss_forward_cuda(
|
|
input,
|
|
target,
|
|
reduction_id,
|
|
n
|
|
)
|
|
|
|
ctx.save_for_backward(input, target)
|
|
ctx.reduction_id = reduction_id
|
|
ctx.N = n
|
|
ctx.op = op
|
|
|
|
return output
|
|
|
|
def backward(ctx, grad_output):
|
|
input, target = ctx.saved_tensors
|
|
|
|
grad_out_scalar = grad_output[0]
|
|
if ctx.reduction_id == 1:
|
|
grad_out_scalar = grad_out_scalar / ctx.N
|
|
|
|
grad_input = torch.empty_like(input)
|
|
grad_target = torch.empty_like(target)
|
|
|
|
ctx.op.generator_loss_backward_cuda(
|
|
grad_out_scalar,
|
|
input,
|
|
target,
|
|
grad_input,
|
|
grad_target,
|
|
ctx.N
|
|
)
|
|
|
|
return grad_input, grad_target, None, None, None
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
def __init__(self, reduction='mean', beta=1.0):
|
|
super().__init__()
|
|
self.beta = float(beta)
|
|
|
|
self.red_map = {'none': 0, 'mean': 1, 'sum': 2}
|
|
if reduction not in self.red_map:
|
|
raise ValueError("Invalid reduction")
|
|
self.reduction_id = self.red_map[reduction]
|
|
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
cpp_source = """
|
|
#include <torch/extension.h>
|
|
#include <vector>
|
|
|
|
std::vector<torch::Tensor> generator_loss_forward_cuda(
|
|
torch::Tensor input,
|
|
torch::Tensor target,
|
|
int reduction,
|
|
int64_t n);
|
|
|
|
void generator_loss_backward_cuda(
|
|
float grad_out_scalar,
|
|
torch::Tensor input,
|
|
torch::Tensor target,
|
|
torch::Tensor grad_input,
|
|
torch::Tensor grad_target,
|
|
int64_t n);
|
|
"""
|
|
|
|
cuda_source = """
|
|
#include <cuda_runtime.h>
|
|
#include <cmath>
|
|
#include <device_functions.h>
|
|
#include <math_functions.h>
|
|
|
|
#define BLOCK_SIZE 256
|
|
#define MAX_GRID_SIZE 4096
|
|
|
|
__inline__ __device__ float warp_reduce_sum(float val) {
|
|
#pragma unroll
|
|
for (int offset = 16; offset > 0; offset /= 2) {
|
|
val += __shfl_down_sync(0xffffffff, val, offset);
|
|
}
|
|
return val;
|
|
}
|
|
|
|
__inline__ __device__ float block_reduce_sum(float val) {
|
|
__shared__ float shared[32];
|
|
int lane = threadIdx.x % 32;
|
|
int wid = threadIdx.x / 32;
|
|
|
|
val = warp_reduce_sum(val);
|
|
if (lane == 0) shared[wid] = val;
|
|
__syncthreads();
|
|
|
|
val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f;
|
|
if (wid == 0) val = warp_reduce_sum(val);
|
|
return val;
|
|
}
|
|
|
|
__global__ void generator_loss_fwd_kernel(
|
|
const float* __restrict__ input,
|
|
const float* __restrict__ target,
|
|
float* __restrict__ output,
|
|
int64_t n,
|
|
int reduction,
|
|
float* __restrict__ partial_sums)
|
|
{
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
int stride = blockDim.x * gridDim.x;
|
|
|
|
float local_sum = 0.0f;
|
|
|
|
for (int idx = i; idx < n; idx += stride) {
|
|
float x = input[idx];
|
|
float t = target[idx];
|
|
|
|
float maxi = (x > 0.0f) ? x : 0.0f;
|
|
float loss = maxi - x * t + __logf(1.0f + __expf(-fabsf(x)));
|
|
|
|
if (reduction == 0) {
|
|
output[idx] = loss;
|
|
} else {
|
|
local_sum += loss;
|
|
}
|
|
}
|
|
|
|
if (reduction != 0) {
|
|
local_sum = block_reduce_sum(local_sum);
|
|
if (threadIdx.x == 0) {
|
|
partial_sums[blockIdx.x] = local_sum;
|
|
}
|
|
}
|
|
}
|
|
|
|
__global__ void reduce_kernel(
|
|
const float* __restrict__ partial_sums,
|
|
float* __restrict__ output,
|
|
int n_partials)
|
|
{
|
|
float local_sum = 0.0f;
|
|
int idx = threadIdx.x;
|
|
|
|
for (int i = idx; i < n_partials; i += blockDim.x) {
|
|
local_sum += partial_sums[i];
|
|
}
|
|
|
|
local_sum = block_reduce_sum(local_sum);
|
|
|
|
if (threadIdx.x == 0) {
|
|
output[0] = local_sum;
|
|
}
|
|
}
|
|
|
|
__global__ void generator_loss_bwd_kernel(
|
|
float grad_out_scalar,
|
|
const float* __restrict__ input,
|
|
const float* __restrict__ target,
|
|
float* __restrict__ grad_input,
|
|
float* __restrict__ grad_target,
|
|
int64_t n)
|
|
{
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
int stride = blockDim.x * gridDim.x;
|
|
|
|
for (int idx = i; idx < n; idx += stride) {
|
|
float x = input[idx];
|
|
float t = target[idx];
|
|
|
|
float sigma_x = 1.0f / (1.0f + __expf(-x));
|
|
|
|
grad_input[idx] = (sigma_x - t) * grad_out_scalar;
|
|
grad_target[idx] = (-x) * grad_out_scalar;
|
|
}
|
|
}
|
|
|
|
std::vector<torch::Tensor> generator_loss_forward_cuda(
|
|
torch::Tensor input,
|
|
torch::Tensor target,
|
|
int reduction,
|
|
int64_t n)
|
|
{
|
|
auto options = input.options();
|
|
torch::Tensor output;
|
|
torch::Tensor partial_sums;
|
|
|
|
const int block_size = BLOCK_SIZE;
|
|
const int grid_size = std::min(
|
|
(int)((n + block_size - 1) / block_size),
|
|
MAX_GRID_SIZE
|
|
);
|
|
|
|
if (reduction == 0) {
|
|
output = torch::empty_like(input);
|
|
partial_sums = torch::empty({1}, options);
|
|
} else {
|
|
output = torch::zeros({1}, options);
|
|
partial_sums = torch::zeros({grid_size}, options);
|
|
}
|
|
|
|
generator_loss_fwd_kernel<<<grid_size, block_size>>>(
|
|
input.data_ptr<float>(),
|
|
target.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
n,
|
|
reduction,
|
|
partial_sums.data_ptr<float>()
|
|
);
|
|
|
|
if (reduction != 0) {
|
|
reduce_kernel<<<1, block_size>>>(
|
|
partial_sums.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
grid_size
|
|
);
|
|
|
|
if (reduction == 1) {
|
|
output.div_(n);
|
|
}
|
|
}
|
|
|
|
return {output, partial_sums};
|
|
}
|
|
|
|
void generator_loss_backward_cuda(
|
|
float grad_out_scalar,
|
|
torch::Tensor input,
|
|
torch::Tensor target,
|
|
torch::Tensor grad_input,
|
|
torch::Tensor grad_target,
|
|
int64_t n)
|
|
{
|
|
const int block_size = BLOCK_SIZE;
|
|
const int grid_size = std::min(
|
|
(int)((n + block_size - 1) / block_size),
|
|
MAX_GRID_SIZE
|
|
);
|
|
|
|
generator_loss_bwd_kernel<<<grid_size, block_size>>>(
|
|
grad_out_scalar,
|
|
input.data_ptr<float>(),
|
|
target.data_ptr<float>(),
|
|
grad_input.data_ptr<float>(),
|
|
grad_target.data_ptr<float>(),
|
|
n
|
|
);
|
|
}
|
|
"""
|
|
|
|
self.op = load_inline(
|
|
name='generator_loss_cuda_v1',
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=['generator_loss_forward_cuda', 'generator_loss_backward_cuda'],
|
|
extra_cuda_cflags=['-O3', '--use_fast_math'],
|
|
verbose=False
|
|
)
|
|
|
|
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
if isinstance(input, (list, tuple)) and len(input) > 0:
|
|
input = input[0]
|
|
target = target[0] if len(target) > 0 else target
|
|
|
|
return GeneratorLossCUDAOp.apply(
|
|
input,
|
|
target,
|
|
self.beta,
|
|
self.reduction_id,
|
|
self.op
|
|
)
|