forked from ccf-ai-infra/GPUCodeForces
126 lines
4.4 KiB
Python
126 lines
4.4 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
cpp_source = """
|
|
#include <torch/extension.h>
|
|
torch::Tensor softmax2d_cuda_forward(torch::Tensor input);
|
|
"""
|
|
|
|
cuda_source = """
|
|
#include <cuda_runtime.h>
|
|
#include <math.h>
|
|
|
|
__global__ void softmax2d_kernel(
|
|
const float* __restrict__ input,
|
|
float* __restrict__ output,
|
|
int channels,
|
|
int spatial_stride, // H * W
|
|
int total_spatial // N * H * W
|
|
) {
|
|
// Flattened spatial index: [0, N*H*W)
|
|
// Maps to (n, h, w) combined
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (idx >= total_spatial) return;
|
|
|
|
// Calculate base pointer offsets
|
|
// Input layout: N, C, H, W
|
|
// We want to iterate over C for a fixed (n, h, w)
|
|
//
|
|
// Decompose idx:
|
|
// n = idx / spatial_stride
|
|
// hw = idx % spatial_stride
|
|
//
|
|
// Address of input[n, c, h, w]:
|
|
// offset = n * (C * HW) + c * HW + hw
|
|
// = c * HW + (n * C * HW + hw)
|
|
// Let base_offset = n * channels * spatial_stride + (idx % spatial_stride)
|
|
|
|
int n = idx / spatial_stride;
|
|
int hw_offset = idx % spatial_stride;
|
|
|
|
// Using long long to prevent overflow for large tensors
|
|
long long base_offset = (long long)n * channels * spatial_stride + hw_offset;
|
|
|
|
// --- Pass 1: Online Max and Sum Calculation ---
|
|
// Reduces global memory reads from 2N (FindMax + CalcSum) to 1N
|
|
|
|
float max_val = -1e37f; // Negative infinity approximation
|
|
float sum_exp = 0.0f;
|
|
|
|
for (int c = 0; c < channels; ++c) {
|
|
// Coalesced read: threads i and i+1 read adjacent memory addresses
|
|
// even though the loop stride is large (H*W).
|
|
long long cur_idx = base_offset + (long long)c * spatial_stride;
|
|
float val = input[cur_idx];
|
|
|
|
if (val > max_val) {
|
|
// Update sum with scaling to prevent overflow
|
|
// sum = sum * exp(old_max - new_max) + exp(new_val - new_max)
|
|
// Note: exp(new_val - new_max) is exp(0) = 1, but we do standard formula
|
|
sum_exp = sum_exp * expf(max_val - val) + 1.0f;
|
|
max_val = val;
|
|
} else {
|
|
sum_exp += expf(val - max_val);
|
|
}
|
|
}
|
|
|
|
// --- Pass 2: Compute Output ---
|
|
// Reciprocal for fast multiplication
|
|
float inv_sum = 1.0f / sum_exp;
|
|
|
|
for (int c = 0; c < channels; ++c) {
|
|
long long cur_idx = base_offset + (long long)c * spatial_stride;
|
|
// Re-read input (L2 cache likely hits if C isn't massive)
|
|
float val = input[cur_idx];
|
|
output[cur_idx] = expf(val - max_val) * inv_sum;
|
|
}
|
|
}
|
|
|
|
torch::Tensor softmax2d_cuda_forward(torch::Tensor input) {
|
|
auto input_contig = input.contiguous();
|
|
|
|
int n = input_contig.size(0);
|
|
int c = input_contig.size(1);
|
|
int h = input_contig.size(2);
|
|
int w = input_contig.size(3);
|
|
|
|
auto output = torch::empty_like(input_contig);
|
|
|
|
int spatial_stride = h * w;
|
|
int total_spatial = n * h * w;
|
|
|
|
int threads = 256;
|
|
int blocks = (total_spatial + threads - 1) / threads;
|
|
|
|
softmax2d_kernel<<<blocks, threads>>>(
|
|
input_contig.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
c,
|
|
spatial_stride,
|
|
total_spatial
|
|
);
|
|
|
|
return output;
|
|
}
|
|
"""
|
|
|
|
self.op = load_inline(
|
|
name="softmax2d_opt",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["softmax2d_cuda_forward"],
|
|
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
|
verbose=False
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.op.softmax2d_cuda_forward(x) |