diff --git a/S1/Ljy123_#64/cudacode.py b/S1/Ljy123_#64/cudacode.py new file mode 100644 index 0000000..74f4b4d --- /dev/null +++ b/S1/Ljy123_#64/cudacode.py @@ -0,0 +1,130 @@ +import torch +from torch.utils.cpp_extension import load_inline + +source = """ +#include +#include +#include + +__device__ __forceinline__ float warp_reduce_sum(float v){ + for(int offset=16; offset>0; offset/=2){ + v += __shfl_down_sync(0xffffffff, v, offset); + } + return v; +} + +__global__ __launch_bounds__(256) void layernorm_silu_kernel(const float* __restrict__ x, const float* __restrict__ gamma, const float* __restrict__ beta, float* __restrict__ y, int B, int D, float eps){ + int b = blockIdx.x; + int tid = threadIdx.x; + int lane = tid & 31; + int warp_id = tid >> 5; + int warps = blockDim.x >> 5; + const float* xr = x + b * D; + float* yr = y + b * D; + __shared__ float s_mean; + __shared__ float s_invstd; + // Welford per-thread + float count = 0.0f; + float mean = 0.0f; + float M2 = 0.0f; + for(int i = tid; i < D; i += blockDim.x){ + float xk = xr[i]; + count += 1.0f; + float delta = xk - mean; + mean += delta / count; + float delta2 = xk - mean; + M2 += delta * delta2; + } + // Warp reduce Welford + for(int offset=16; offset>0; offset/=2){ + float mean_other = __shfl_down_sync(0xffffffff, mean, offset); + float M2_other = __shfl_down_sync(0xffffffff, M2, offset); + float cnt_other = __shfl_down_sync(0xffffffff, count, offset); + if(cnt_other > 0.0f){ + float delta = mean_other - mean; + float tot = count + cnt_other; + M2 += M2_other + delta * delta * (count * cnt_other) / tot; + mean += (mean_other - mean) * (cnt_other / tot); + count = tot; + } + } + __shared__ float s_mean_warp[32]; + __shared__ float s_M2_warp[32]; + __shared__ float s_cnt_warp[32]; + if(lane == 0){ + s_mean_warp[warp_id] = mean; + s_M2_warp[warp_id] = M2; + s_cnt_warp[warp_id] = count; + } + __syncthreads(); + if(warp_id == 0){ + float mean0 = (lane < warps) ? s_mean_warp[lane] : 0.0f; + float M20 = (lane < warps) ? s_M2_warp[lane] : 0.0f; + float cnt0 = (lane < warps) ? s_cnt_warp[lane] : 0.0f; + // reduce across warps + for(int offset=16; offset>0; offset/=2){ + float mean_other = __shfl_down_sync(0xffffffff, mean0, offset); + float M2_other = __shfl_down_sync(0xffffffff, M20, offset); + float cnt_other = __shfl_down_sync(0xffffffff, cnt0, offset); + if(cnt_other > 0.0f){ + float delta = mean_other - mean0; + float tot = cnt0 + cnt_other; + M20 += M2_other + delta * delta * (cnt0 * cnt_other) / tot; + mean0 += (mean_other - mean0) * (cnt_other / tot); + cnt0 = tot; + } + } + if(lane == 0){ + float var = M20 / fmaxf(cnt0, 1.0f); + s_mean = mean0; + s_invstd = rsqrtf(fmaxf(var, 0.0f) + eps); + } + } + __syncthreads(); + for(int i = tid; i < D; i += blockDim.x){ + float v = (__ldg(xr + i) - s_mean) * s_invstd; + float u = v * __ldg(gamma + i) + __ldg(beta + i); + float g = __fdividef(1.0f, 1.0f + __expf(-u)); + yr[i] = u * g; + } +} + +torch::Tensor layernorm_silu_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor eps){ + auto xc = x.contiguous(); + auto gc = gamma.contiguous(); + auto bc = beta.contiguous(); + auto y = torch::empty_like(xc); + int B = (int)xc.size(0); + int D = (int)xc.size(1); + float e = eps.item(); + int block = 256; + dim3 grid(B); + layernorm_silu_kernel<<>>(xc.data_ptr(), gc.data_ptr(), bc.data_ptr(), y.data_ptr(), B, D, e); + return y; +} +""" + +cpp_source = """ +torch::Tensor layernorm_silu_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor eps); +""" + +ops = load_inline( + name="layernorm_silu", + cpp_sources=cpp_source, + cuda_sources=source, + functions=["layernorm_silu_cuda"], + extra_cflags=["-O3","-std=c++17"], + extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17"], + verbose=True +) + +class ModelNew(torch.nn.Module): + def __init__(self, gamma: torch.Tensor, beta: torch.Tensor, eps: float): + super(ModelNew, self).__init__() + self.ops = ops + self.register_buffer("gamma", gamma) + self.register_buffer("beta", beta) + self.register_buffer("eps", torch.tensor(float(eps), dtype=torch.float32)) + + def forward(self, x): + return self.ops.layernorm_silu_cuda(x, self.gamma, self.beta, self.eps) diff --git a/S1/Ljy123_#64/prompt.txt b/S1/Ljy123_#64/prompt.txt new file mode 100644 index 0000000..3426644 --- /dev/null +++ b/S1/Ljy123_#64/prompt.txt @@ -0,0 +1,15 @@ +Objective: Fused LayerNorm + SiLU CUDA kernel with rtol=1e-3 accuracy and ≥1.3x speedup. + +Computation: +- Row layer normalization (mean/var over dim) with epsilon, then affine transform, then SiLU: `y = u * sigmoid(u)`. + +Numerical Method: +- Use Welford’s online algorithm for mean/variance to match PyTorch precision and avoid catastrophic cancellation. +- Perform warp-level reductions and shared-memory aggregation across warps. + +Performance: +- 512 threads per block; coalesced loads and parallel write-back. +- Fuse normalization, affine, and activation to minimize memory traffic. + +Benchmark: +- Batch size: 16; Dim: 16384; 100 iterations; target speedup ≥1.3x. diff --git a/S1/Ljy123_#64/run_code.py b/S1/Ljy123_#64/run_code.py new file mode 100644 index 0000000..4aac3b0 --- /dev/null +++ b/S1/Ljy123_#64/run_code.py @@ -0,0 +1,51 @@ +import torch +import time +from torchcode import Model, get_inputs, get_init_inputs +from cudacode import ModelNew + +def run_benchmark(): + if not torch.cuda.is_available(): + print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。") + return + device = torch.device("cuda") + + init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()] + inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()] + + torch_model = Model(*init_inputs).cuda().eval() + cuda_model = ModelNew(*init_inputs).cuda().eval() + + print("-------------------- 精度对齐验证 --------------------") + with torch.no_grad(): + output_torch = torch_model(*inputs) + output_cuda = cuda_model(*inputs) + precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03) + if precision_flag: + print("✅ 精度对齐:两个模型的输出结果非常接近。") + else: + print("❌ 精度不一致!") + + print("\n-------------------- 性能加速比测试 --------------------") + num_iterations = 100 + torch.cuda.synchronize(); start_time = time.time() + for _ in range(num_iterations): + _ = torch_model(*inputs) + torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations + + torch.cuda.synchronize(); start_time = time.time() + for _ in range(num_iterations): + _ = cuda_model(*inputs) + torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations + + print(f"PyTorch LayerNorm-SiLU 平均执行时间: {torch_time:.6f} 秒") + print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f} 秒") + speedup = torch_time / cuda_time if cuda_time > 0 else 0 + if cuda_time > 0: + print(f"加速比 (Speedup): {speedup:.2f}x") + else: + print("CUDA 内核执行时间为0,无法计算加速比。") + return precision_flag, speedup + +if __name__ == "__main__": + run_benchmark() + diff --git a/S1/Ljy123_#64/torchcode.py b/S1/Ljy123_#64/torchcode.py new file mode 100644 index 0000000..c701f35 --- /dev/null +++ b/S1/Ljy123_#64/torchcode.py @@ -0,0 +1,30 @@ +import torch +import torch.nn as nn + +class Model(nn.Module): + def __init__(self, gamma: torch.Tensor, beta: torch.Tensor, eps: float): + super(Model, self).__init__() + self.register_buffer("gamma", gamma) + self.register_buffer("beta", beta) + self.register_buffer("eps", torch.tensor(float(eps), dtype=torch.float32)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + mean = x.mean(dim=1, keepdim=True) + var = x.var(dim=1, unbiased=False, keepdim=True) + xn = (x - mean) / torch.sqrt(var + self.eps) + u = xn * self.gamma + self.beta + g = torch.sigmoid(u) + return u * g + +batch_size = 16 +dim = 16384 + +def get_inputs(): + x = torch.randn(batch_size, dim) + return [x] + +def get_init_inputs(): + gamma = torch.randn(dim) + beta = torch.randn(dim) + return [gamma, beta, 1e-5] +