diff --git a/S1/hli28146_#36/conjphysical_cuda.py b/S1/hli28146_#36/conjphysical_cuda.py new file mode 100644 index 00000000..e61c09f4 --- /dev/null +++ b/S1/hli28146_#36/conjphysical_cuda.py @@ -0,0 +1,111 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline + +cpp_source = """ +#include + +torch::Tensor conj_physical_cuda_forward(const torch::Tensor& input); +""" + +cuda_source = """ +#include +#include + +struct __align__(16) Float4 { + float x, y, z, w; // R1, I1, R2, I2 +}; + +struct __align__(8) Float2 { + float x, y; // R, I +}; + +// Core logic: z = x + iy -> z* = x - iy +// We manipulate raw floats to avoid complex class overhead +__global__ void conj_physical_kernel( + const float* __restrict__ input, + float* __restrict__ output, + const int n_complex_elements) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + int vec_loops = n_complex_elements / 2; + + const Float4* vec_input = reinterpret_cast(input); + Float4* vec_output = reinterpret_cast(output); + + for (int i = idx; i < vec_loops; i += stride) { + Float4 val = vec_input[i]; + + // Logical layout: x=Real1, y=Imag1, z=Real2, w=Imag2 + // Operation: Negate Imag parts + val.y = -val.y; + val.w = -val.w; + + vec_output[i] = val; + } + + // Only happens if n_complex_elements is odd + int tail_idx = vec_loops * 2; + + // We switch to Float2 pointer to access single complex elements + const Float2* scalar_input = reinterpret_cast(input); + Float2* scalar_output = reinterpret_cast(output); + + // Standard grid stride logic applied to the tail part + // Though usually this loop runs at most once per thread if aligned + for (int i = tail_idx + idx; i < n_complex_elements; i += stride) { + Float2 val = scalar_input[i]; + val.y = -val.y; // Negate Imag + scalar_output[i] = val; + } +} + +torch::Tensor conj_physical_cuda_forward(const torch::Tensor& input) { + TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor"); + TORCH_CHECK(input.is_contiguous(), "Input must be contiguous"); + TORCH_CHECK(input.scalar_type() == torch::kComplexFloat, "Input must be ComplexFloat (complex64)"); + + int n_elements = input.numel(); + auto output = torch::empty_like(input); + + const int block_size = 256; + // Each thread handles 2 elements ideally + int num_vectors = (n_elements + 1) / 2; + int grid_size = (num_vectors + block_size - 1) / block_size; + if (grid_size > 65535) grid_size = 65535; + + conj_physical_kernel<<>>( + reinterpret_cast(input.data_ptr>()), + reinterpret_cast(output.data_ptr>()), + n_elements + ); + + return output; +} +""" + +conj_op = load_inline( + name='conj_physical_op', + cpp_sources=cpp_source, + cuda_sources=cuda_source, + functions=['conj_physical_cuda_forward'], + verbose=False, + extra_cuda_cflags=['-O3'] +) + +class ConjNew(nn.Module): + def __init__(self): + super(ConjNew, self).__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return conj_op.conj_physical_cuda_forward(x) + +class ModelNew(nn.Module): + def __init__(self): + super(ModelNew, self).__init__() + self.act = ConjNew() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.act(x) \ No newline at end of file diff --git a/S1/hli28146_#36/conjphysical_torch.py b/S1/hli28146_#36/conjphysical_torch.py new file mode 100644 index 00000000..340b4a7f --- /dev/null +++ b/S1/hli28146_#36/conjphysical_torch.py @@ -0,0 +1,28 @@ +import torch +import torch.nn as nn + +BATCH_SIZE = 4096 +DIM = 4096 +SHAPE = (BATCH_SIZE, DIM) + +class ConjPhysicalModel(nn.Module): + def __init__(self): + super(ConjPhysicalModel, self).__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.conj_physical(x) + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.act = ConjPhysicalModel() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.act(x) + +def get_inputs(): + x = torch.randn(SHAPE, dtype=torch.complex64) + return [x.contiguous()] + +def get_init_inputs(): + return [] \ No newline at end of file diff --git a/S1/hli28146_#36/prompt.txt b/S1/hli28146_#36/prompt.txt new file mode 100644 index 00000000..ceb57422 --- /dev/null +++ b/S1/hli28146_#36/prompt.txt @@ -0,0 +1,54 @@ +Write a custom CUDA kernel to optimize `torch.conj_physical` for complex tensors. + +The operation computes the element-wise conjugate of a complex tensor. For z = x + iy, conj_physical(z) = x - iy. It explicitly materializes the result in memory. + +Problem Analysis: +This is a strictly memory-bound operation. +1. Data Layout: `complex64` stores data as contiguous pairs of floats [Real, Imag]. +2. Computation: The only arithmetic operation is negating the imaginary part. +3. Bottleneck: The performance is strictly limited by Global Memory bandwidth. + +Optimization Strategy: Vectorized Access (2x Complex Elements per Thread) + +1. Vectorized I/O (Float4): + - A single `complex64` is 8 bytes (2 floats). + - Using `float4` (16 bytes) allows a single thread to load/store **two** complex numbers at once. + - Layout loaded into registers: `x`=Real1, `y`=Imag1, `z`=Real2, `w`=Imag2. + +2. In-Register Computation: + - Negate the `y` and `w` components (the imaginary parts). + - Store the modified `float4` back to global memory. + +3. Grid-Stride Loop: Implement a robust grid-stride loop to handle arbitrary tensor sizes, processing 2 complex elements per iteration in the vectorized loop, and handling remainders with a scalar loop. + +Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: + +```python +import torch +import torch.nn as nn + +BATCH_SIZE = 4096 +DIM = 4096 +SHAPE = (BATCH_SIZE, DIM) + +class ConjPhysicalModel(nn.Module): + def __init__(self): + super(ConjPhysicalModel, self).__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.conj_physical(x) + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.act = ConjPhysicalModel() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.act(x) + +def get_inputs(): + x = torch.randn(SHAPE, dtype=torch.complex64) + return [x.contiguous()] + +def get_init_inputs(): + return [] \ No newline at end of file diff --git a/S1/hli28146_#36/run_code.py b/S1/hli28146_#36/run_code.py new file mode 100644 index 00000000..8b99d6cd --- /dev/null +++ b/S1/hli28146_#36/run_code.py @@ -0,0 +1,74 @@ +########################################################### +# 性能和精度验证程序 +########################################################### +import torch +import torch.nn as nn +import time +from conjphysical_torch import Model,get_inputs,get_init_inputs +from conjphysical_cuda import ModelNew + +def run_benchmark(): + # 检查 CUDA 是否可用 + if not torch.cuda.is_available(): + print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。") + return + else: + device = torch.device("cuda") + + # 初始化模型 + init_inputs = get_init_inputs() + init_inputs = [ + x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs + ] + inputs = get_inputs() + inputs = [ + x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs + ] + + torch_model = Model(*init_inputs).cuda() + cuda_model = ModelNew(*init_inputs).cuda() + + torch_model.eval() + cuda_model.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 + + # PyTorch 模型计时 + 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 + + # 自定义 CUDA 内核计时 + 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 torch.relu 平均执行时间: {torch_time:.6f} 秒") + print(f"自定义 CUDA 内核 平均执行时间: {cuda_time:.6f} 秒") + speedup = 0 + if cuda_time > 0: + speedup = torch_time / cuda_time + print(f"加速比 (Speedup): {speedup:.2f}x") + else: + print("CUDA 内核执行时间为0,无法计算加速比。") + return precision_flag,speedup +if __name__ == "__main__": + precision_flag,speedup = run_benchmark() \ No newline at end of file