From bb6f723ef51d2bb474ce938910cebf2b50f2eb67 Mon Sep 17 00:00:00 2001 From: gsd <2396583337@qq.com> Date: Tue, 9 Dec 2025 11:50:28 +0800 Subject: [PATCH 1/2] finish segment_max #130 --- S1/gsd123_#130/prompt.txt | 92 +++++++++++++++++++++++++++ S1/gsd123_#130/run_code.py | 77 +++++++++++++++++++++++ S1/gsd123_#130/segmentmax_cuda.py | 99 ++++++++++++++++++++++++++++++ S1/gsd123_#130/segmentmax_torch.py | 25 ++++++++ 4 files changed, 293 insertions(+) create mode 100644 S1/gsd123_#130/prompt.txt create mode 100644 S1/gsd123_#130/run_code.py create mode 100644 S1/gsd123_#130/segmentmax_cuda.py create mode 100644 S1/gsd123_#130/segmentmax_torch.py diff --git a/S1/gsd123_#130/prompt.txt b/S1/gsd123_#130/prompt.txt new file mode 100644 index 0000000..db1eb1f --- /dev/null +++ b/S1/gsd123_#130/prompt.txt @@ -0,0 +1,92 @@ +You write custom CUDA kernels to replace the pytorch operators in the given GeGLU architecture to get speedups. + +You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining chunk+gelu+elementwise_mul), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination. + +PyTorch C++/CUDA Extension + +Runtime compilation via torch.utils.cpp_extension.load_inline + +Direct integration with PyTorch's autograd and tensor ecosystem + +CUDA Custom Atomic Operation + +Custom atomicMaxFloat using atomicCAS (Compare-And-Swap) + +Handles float atomic maximum (CUDA lacks native atomicMax for floats) + +Uses __float_as_int/__int_as_float for type-punning atomic operations + +Implements spin-loop with early exit for performance + +Two-Kernel Design + +init_kernel: Initializes output to -1e38 (negative infinity proxy) + +segment_max_kernel: Performs segment-wise maximum reduction + +CUDA Optimization Techniques + +__forceinline__ for device function inlining + +__restrict__ pointers for compiler aliasing optimization + +Grid-stride loops with 256 threads per block + +Efficient 2D→1D index calculation: row * channels + col + +Numerical Stability Pattern + +Initializes with -1e38f instead of -FLT_MAX for safety + +Uses fmaxf for maximum computation (CUDA math intrinsic) + +Memory Management + +torch::empty for uninitialized tensor creation (faster than zeros) + +Separate initialization kernel for output tensor + +Coalesced global memory access patterns + +Segment-Based Reduction + +Groups input rows by segment_id from index tensor + +Computes channel-wise maximum within each segment + +Output shape: (dim_size, channels) + +Module Abstraction + +nn.Module wrapper for PyTorch integration + +Maintains dim_size as persistent configuration + + + +Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: +import torch +import torch.nn as nn + +class Model(nn.Module): + def __init__(self, dim_size): + super(Model, self).__init__() + self.dim_size = dim_size + + def forward(self, src, index): + out = torch.full((self.dim_size, src.size(1)), -1e38, device=src.device, dtype=src.dtype) + index_expanded = index.unsqueeze(1).expand_as(src) + out.scatter_reduce_(0, index_expanded, src, reduce='amax', include_self=False) + return out + +batch_size = 1024 +features = 64 +dim_size = 128 + +def get_inputs(): + src = torch.randn(batch_size, features) + index = torch.randint(0, dim_size, (batch_size,)) + return [src, index] + +def get_init_inputs(): + return [dim_s \ No newline at end of file diff --git a/S1/gsd123_#130/run_code.py b/S1/gsd123_#130/run_code.py new file mode 100644 index 0000000..e0e5aa9 --- /dev/null +++ b/S1/gsd123_#130/run_code.py @@ -0,0 +1,77 @@ +########################################################### +# 性能和精度验证程序 +########################################################### +import torch +import torch.nn as nn +import time +from segmentmax_torch import Model, get_inputs, get_init_inputs +from segmentmax_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 diff --git a/S1/gsd123_#130/segmentmax_cuda.py b/S1/gsd123_#130/segmentmax_cuda.py new file mode 100644 index 0000000..755d89b --- /dev/null +++ b/S1/gsd123_#130/segmentmax_cuda.py @@ -0,0 +1,99 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline + +cuda_source = """ +#include +#include +#include + +__device__ __forceinline__ void atomicMaxFloat(float* address, float val) { + int* address_as_i = (int*)address; + int old = *address_as_i, assumed; + do { + assumed = old; + float old_val = __int_as_float(assumed); + float new_val = fmaxf(val, old_val); + if (new_val == old_val) break; + old = atomicCAS(address_as_i, assumed, __float_as_int(new_val)); + } while (assumed != old); +} + +__global__ void init_kernel(float* out, int size, float val) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < size) { + out[idx] = val; + } +} + +__global__ void segment_max_kernel( + const float* __restrict__ src, + const long* __restrict__ index, + float* __restrict__ out, + int num_elements, + int channels, + int dim_size +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_elements) { + int row = idx / channels; + int col = idx % channels; + + long segment_id = index[row]; + + if (segment_id >= 0 && segment_id < dim_size) { + int out_idx = segment_id * channels + col; + atomicMaxFloat(&out[out_idx], src[idx]); + } + } +} + +torch::Tensor segment_max_cuda(torch::Tensor src, torch::Tensor index, int dim_size) { + int N = src.size(0); + int C = src.size(1); + + auto out = torch::empty({dim_size, C}, src.options()); + + int out_elements = dim_size * C; + const int block_size = 256; + int grid_init = (out_elements + block_size - 1) / block_size; + + init_kernel<<>>(out.data_ptr(), out_elements, -1e38f); + + int num_elements = N * C; + int grid_max = (num_elements + block_size - 1) / block_size; + + segment_max_kernel<<>>( + src.data_ptr(), + index.data_ptr(), + out.data_ptr(), + num_elements, + C, + dim_size + ); + + return out; +} +""" + +cpp_source = """ +torch::Tensor segment_max_cuda(torch::Tensor src, torch::Tensor index, int dim_size); +""" + +segment_max_lib = load_inline( + name="segment_max", + cpp_sources=cpp_source, + cuda_sources=cuda_source, + functions=["segment_max_cuda"], + verbose=True +) + + +class ModelNew(nn.Module): + def __init__(self, dim_size): + super(ModelNew, self).__init__() + self.dim_size = dim_size + self.lib = segment_max_lib + + def forward(self, src, index): + return self.lib.segment_max_cuda(src, index, self.dim_size) \ No newline at end of file diff --git a/S1/gsd123_#130/segmentmax_torch.py b/S1/gsd123_#130/segmentmax_torch.py new file mode 100644 index 0000000..af10627 --- /dev/null +++ b/S1/gsd123_#130/segmentmax_torch.py @@ -0,0 +1,25 @@ +import torch +import torch.nn as nn + +class Model(nn.Module): + def __init__(self, dim_size): + super(Model, self).__init__() + self.dim_size = dim_size + + def forward(self, src, index): + out = torch.full((self.dim_size, src.size(1)), -1e38, device=src.device, dtype=src.dtype) + index_expanded = index.unsqueeze(1).expand_as(src) + out.scatter_reduce_(0, index_expanded, src, reduce='amax', include_self=False) + return out + +batch_size = 1024 +features = 64 +dim_size = 128 + +def get_inputs(): + src = torch.randn(batch_size, features) + index = torch.randint(0, dim_size, (batch_size,)) + return [src, index] + +def get_init_inputs(): + return [dim_size] \ No newline at end of file From f2e8418c9a8cdd7f879721828151c1e1e21d8bd5 Mon Sep 17 00:00:00 2001 From: gsd <2396583337@qq.com> Date: Tue, 9 Dec 2025 15:11:39 +0800 Subject: [PATCH 2/2] finish SSIMLoss #130 --- S1/gsd123_#130/SSIMLoss_cuda.py | 128 +++++++++++++++++++++++++++++ S1/gsd123_#130/SSIMLoss_torch.py | 61 ++++++++++++++ S1/gsd123_#130/prompt.txt | 124 ++++++++++++++-------------- S1/gsd123_#130/run_code.py | 4 +- S1/gsd123_#130/segmentmax_cuda.py | 99 ---------------------- S1/gsd123_#130/segmentmax_torch.py | 25 ------ 6 files changed, 252 insertions(+), 189 deletions(-) create mode 100644 S1/gsd123_#130/SSIMLoss_cuda.py create mode 100644 S1/gsd123_#130/SSIMLoss_torch.py delete mode 100644 S1/gsd123_#130/segmentmax_cuda.py delete mode 100644 S1/gsd123_#130/segmentmax_torch.py diff --git a/S1/gsd123_#130/SSIMLoss_cuda.py b/S1/gsd123_#130/SSIMLoss_cuda.py new file mode 100644 index 0000000..1625443 --- /dev/null +++ b/S1/gsd123_#130/SSIMLoss_cuda.py @@ -0,0 +1,128 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline +import math + +cuda_source = """ +#include +#include + +#define WINDOW_SIZE 11 +#define RADIUS 5 +#define C1 (0.01f * 0.01f) +#define C2 (0.03f * 0.03f) + +__global__ void fused_ssim_kernel( + const float* __restrict__ img1, + const float* __restrict__ img2, + const float* __restrict__ gaussian_kernel, + float* __restrict__ out_map, + int B, int C, int H, int W +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total_pixels = B * C * H * W; + + if (idx >= total_pixels) return; + + int w_idx = idx % W; + int h_idx = (idx / W) % H; + int c_idx = (idx / (W * H)) % C; + int b_idx = idx / (W * H * C); + + int pixel_offset = b_idx * (C * H * W) + c_idx * (H * W); + + float mu1 = 0.0f; + float mu2 = 0.0f; + float sigma1_sq_sum = 0.0f; + float sigma2_sq_sum = 0.0f; + float sigma12_sum = 0.0f; + + for (int i = -RADIUS; i <= RADIUS; ++i) { + for (int j = -RADIUS; j <= RADIUS; ++j) { + int cur_h = h_idx + i; + int cur_w = w_idx + j; + + float val1 = 0.0f; + float val2 = 0.0f; + + if (cur_h >= 0 && cur_h < H && cur_w >= 0 && cur_w < W) { + int neighbor_idx = pixel_offset + cur_h * W + cur_w; + val1 = img1[neighbor_idx]; + val2 = img2[neighbor_idx]; + } + + float weight = gaussian_kernel[(i + RADIUS) * WINDOW_SIZE + (j + RADIUS)]; + + mu1 += weight * val1; + mu2 += weight * val2; + sigma1_sq_sum += weight * val1 * val1; + sigma2_sq_sum += weight * val2 * val2; + sigma12_sum += weight * val1 * val2; + } + } + + float mu1_sq = mu1 * mu1; + float mu2_sq = mu2 * mu2; + float mu1_mu2 = mu1 * mu2; + + float sigma1_sq = sigma1_sq_sum - mu1_sq; + float sigma2_sq = sigma2_sq_sum - mu2_sq; + float sigma12 = sigma12_sum - mu1_mu2; + + float num = (2.0f * mu1_mu2 + C1) * (2.0f * sigma12 + C2); + float den = (mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2); + + out_map[idx] = num / den; +} + +torch::Tensor ssim_cuda(torch::Tensor img1, torch::Tensor img2, torch::Tensor kernel) { + int B = img1.size(0); + int C = img1.size(1); + int H = img1.size(2); + int W = img1.size(3); + + auto out_map = torch::empty_like(img1); + + int total_pixels = B * C * H * W; + int threads = 256; + int blocks = (total_pixels + threads - 1) / threads; + + fused_ssim_kernel<<>>( + img1.data_ptr(), + img2.data_ptr(), + kernel.data_ptr(), + out_map.data_ptr(), + B, C, H, W + ); + + return 1.0f - out_map.mean(); +} +""" + +cpp_source = """ +torch::Tensor ssim_cuda(torch::Tensor img1, torch::Tensor img2, torch::Tensor kernel); +""" + +ssim_loss = load_inline( + name="ssim_loss", + cpp_sources=cpp_source, + cuda_sources=cuda_source, + functions=["ssim_cuda"], + verbose=False +) + + +class ModelNew(nn.Module): + def __init__(self, window_size=11, sigma=1.5): + super(ModelNew, self).__init__() + self.register_buffer("kernel", self._create_kernel(window_size, sigma)) + + def _create_kernel(self, window_size, sigma): + coords = torch.arange(window_size).float() - window_size // 2 + g = torch.exp(-(coords ** 2) / (2 * sigma ** 2)) + g = g / g.sum() + kernel = g.unsqueeze(1) @ g.unsqueeze(0) + return kernel.contiguous() + + def forward(self, img1, img2): + return ssim_loss.ssim_cuda(img1, img2, self.kernel) \ No newline at end of file diff --git a/S1/gsd123_#130/SSIMLoss_torch.py b/S1/gsd123_#130/SSIMLoss_torch.py new file mode 100644 index 0000000..f689d17 --- /dev/null +++ b/S1/gsd123_#130/SSIMLoss_torch.py @@ -0,0 +1,61 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import math + + +class Model(nn.Module): + def __init__(self, window_size=11, channel=3): + super(Model, self).__init__() + self.window_size = window_size + self.channel = channel + self.window = self.create_window(window_size, channel) + + def create_window(self, window_size, channel): + def _gaussian(window_size, sigma): + gauss = torch.Tensor( + [math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + + _1D_window = _gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = _2D_window.expand(channel, 1, window_size, window_size).contiguous() + return window + + def forward(self, img1, img2): + if self.window.device != img1.device: + self.window = self.window.to(img1.device) + self.window = self.window.type_as(img1) + + mu1 = F.conv2d(img1, self.window, padding=self.window_size // 2, groups=self.channel) + mu2 = F.conv2d(img2, self.window, padding=self.window_size // 2, groups=self.channel) + + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + + sigma1_sq = F.conv2d(img1 * img1, self.window, padding=self.window_size // 2, groups=self.channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, self.window, padding=self.window_size // 2, groups=self.channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, self.window, padding=self.window_size // 2, groups=self.channel) - mu1_mu2 + + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + + ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return 1 - ssim_map.mean() + + +batch_size = 16 +channels = 3 +height = 256 +width = 256 + + +def get_inputs(): + img1 = torch.rand(batch_size, channels, height, width) + img2 = torch.rand(batch_size, channels, height, width) + return [img1, img2] + + +def get_init_inputs(): + return [] \ No newline at end of file diff --git a/S1/gsd123_#130/prompt.txt b/S1/gsd123_#130/prompt.txt index db1eb1f..e024f66 100644 --- a/S1/gsd123_#130/prompt.txt +++ b/S1/gsd123_#130/prompt.txt @@ -1,92 +1,90 @@ You write custom CUDA kernels to replace the pytorch operators in the given GeGLU architecture to get speedups. You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining chunk+gelu+elementwise_mul), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination. +Custom CUDA kernel extension via torch.utils.cpp_extension.load_inline -PyTorch C++/CUDA Extension +Structural Similarity Index (SSIM) loss computation -Runtime compilation via torch.utils.cpp_extension.load_inline +Fused window-based SSIM calculation (11×11 Gaussian window) -Direct integration with PyTorch's autograd and tensor ecosystem +Local statistics computation: means, variances, covariance -CUDA Custom Atomic Operation +Gaussian kernel weighting for spatial weighting -Custom atomicMaxFloat using atomicCAS (Compare-And-Swap) +SSIM formula: (2μ₁μ₂ + C₁)(2σ₁₂ + C₂) / ((μ₁² + μ₂² + C₁)(σ₁² + σ₂² + C₂)) -Handles float atomic maximum (CUDA lacks native atomicMax for floats) +Boundary handling with conditional checks -Uses __float_as_int/__int_as_float for type-punning atomic operations +Element-wise parallelization across all pixels×channels×batches -Implements spin-loop with early exit for performance +Fixed block size (256 threads) with dynamic grid sizing -Two-Kernel Design +Mean reduction across all pixels (1 - SSIM mean) -init_kernel: Initializes output to -1e38 (negative infinity proxy) +Precomputed Gaussian kernel as PyTorch buffer -segment_max_kernel: Performs segment-wise maximum reduction - -CUDA Optimization Techniques - -__forceinline__ for device function inlining - -__restrict__ pointers for compiler aliasing optimization - -Grid-stride loops with 256 threads per block - -Efficient 2D→1D index calculation: row * channels + col - -Numerical Stability Pattern - -Initializes with -1e38f instead of -FLT_MAX for safety - -Uses fmaxf for maximum computation (CUDA math intrinsic) - -Memory Management - -torch::empty for uninitialized tensor creation (faster than zeros) - -Separate initialization kernel for output tensor - -Coalesced global memory access patterns - -Segment-Based Reduction - -Groups input rows by segment_id from index tensor - -Computes channel-wise maximum within each segment - -Output shape: (dim_size, channels) - -Module Abstraction - -nn.Module wrapper for PyTorch integration - -Maintains dim_size as persistent configuration Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: import torch import torch.nn as nn +import torch.nn.functional as F +import math + class Model(nn.Module): - def __init__(self, dim_size): + def __init__(self, window_size=11, channel=3): super(Model, self).__init__() - self.dim_size = dim_size + self.window_size = window_size + self.channel = channel + self.window = self.create_window(window_size, channel) - def forward(self, src, index): - out = torch.full((self.dim_size, src.size(1)), -1e38, device=src.device, dtype=src.dtype) - index_expanded = index.unsqueeze(1).expand_as(src) - out.scatter_reduce_(0, index_expanded, src, reduce='amax', include_self=False) - return out + def create_window(self, window_size, channel): + def _gaussian(window_size, sigma): + gauss = torch.Tensor( + [math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + + _1D_window = _gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = _2D_window.expand(channel, 1, window_size, window_size).contiguous() + return window + + def forward(self, img1, img2): + if self.window.device != img1.device: + self.window = self.window.to(img1.device) + self.window = self.window.type_as(img1) + + mu1 = F.conv2d(img1, self.window, padding=self.window_size // 2, groups=self.channel) + mu2 = F.conv2d(img2, self.window, padding=self.window_size // 2, groups=self.channel) + + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + + sigma1_sq = F.conv2d(img1 * img1, self.window, padding=self.window_size // 2, groups=self.channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, self.window, padding=self.window_size // 2, groups=self.channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, self.window, padding=self.window_size // 2, groups=self.channel) - mu1_mu2 + + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + + ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return 1 - ssim_map.mean() + + +batch_size = 16 +channels = 3 +height = 256 +width = 256 -batch_size = 1024 -features = 64 -dim_size = 128 def get_inputs(): - src = torch.randn(batch_size, features) - index = torch.randint(0, dim_size, (batch_size,)) - return [src, index] + img1 = torch.rand(batch_size, channels, height, width) + img2 = torch.rand(batch_size, channels, height, width) + return [img1, img2] + def get_init_inputs(): - return [dim_s \ No newline at end of file + return [] \ No newline at end of file diff --git a/S1/gsd123_#130/run_code.py b/S1/gsd123_#130/run_code.py index e0e5aa9..b6011a3 100644 --- a/S1/gsd123_#130/run_code.py +++ b/S1/gsd123_#130/run_code.py @@ -4,8 +4,8 @@ import torch import torch.nn as nn import time -from segmentmax_torch import Model, get_inputs, get_init_inputs -from segmentmax_cuda import ModelNew +from SSIMLoss_torch import Model, get_inputs, get_init_inputs +from SSIMLoss_cuda import ModelNew def run_benchmark(): diff --git a/S1/gsd123_#130/segmentmax_cuda.py b/S1/gsd123_#130/segmentmax_cuda.py deleted file mode 100644 index 755d89b..0000000 --- a/S1/gsd123_#130/segmentmax_cuda.py +++ /dev/null @@ -1,99 +0,0 @@ -import torch -import torch.nn as nn -from torch.utils.cpp_extension import load_inline - -cuda_source = """ -#include -#include -#include - -__device__ __forceinline__ void atomicMaxFloat(float* address, float val) { - int* address_as_i = (int*)address; - int old = *address_as_i, assumed; - do { - assumed = old; - float old_val = __int_as_float(assumed); - float new_val = fmaxf(val, old_val); - if (new_val == old_val) break; - old = atomicCAS(address_as_i, assumed, __float_as_int(new_val)); - } while (assumed != old); -} - -__global__ void init_kernel(float* out, int size, float val) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx < size) { - out[idx] = val; - } -} - -__global__ void segment_max_kernel( - const float* __restrict__ src, - const long* __restrict__ index, - float* __restrict__ out, - int num_elements, - int channels, - int dim_size -) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx < num_elements) { - int row = idx / channels; - int col = idx % channels; - - long segment_id = index[row]; - - if (segment_id >= 0 && segment_id < dim_size) { - int out_idx = segment_id * channels + col; - atomicMaxFloat(&out[out_idx], src[idx]); - } - } -} - -torch::Tensor segment_max_cuda(torch::Tensor src, torch::Tensor index, int dim_size) { - int N = src.size(0); - int C = src.size(1); - - auto out = torch::empty({dim_size, C}, src.options()); - - int out_elements = dim_size * C; - const int block_size = 256; - int grid_init = (out_elements + block_size - 1) / block_size; - - init_kernel<<>>(out.data_ptr(), out_elements, -1e38f); - - int num_elements = N * C; - int grid_max = (num_elements + block_size - 1) / block_size; - - segment_max_kernel<<>>( - src.data_ptr(), - index.data_ptr(), - out.data_ptr(), - num_elements, - C, - dim_size - ); - - return out; -} -""" - -cpp_source = """ -torch::Tensor segment_max_cuda(torch::Tensor src, torch::Tensor index, int dim_size); -""" - -segment_max_lib = load_inline( - name="segment_max", - cpp_sources=cpp_source, - cuda_sources=cuda_source, - functions=["segment_max_cuda"], - verbose=True -) - - -class ModelNew(nn.Module): - def __init__(self, dim_size): - super(ModelNew, self).__init__() - self.dim_size = dim_size - self.lib = segment_max_lib - - def forward(self, src, index): - return self.lib.segment_max_cuda(src, index, self.dim_size) \ No newline at end of file diff --git a/S1/gsd123_#130/segmentmax_torch.py b/S1/gsd123_#130/segmentmax_torch.py deleted file mode 100644 index af10627..0000000 --- a/S1/gsd123_#130/segmentmax_torch.py +++ /dev/null @@ -1,25 +0,0 @@ -import torch -import torch.nn as nn - -class Model(nn.Module): - def __init__(self, dim_size): - super(Model, self).__init__() - self.dim_size = dim_size - - def forward(self, src, index): - out = torch.full((self.dim_size, src.size(1)), -1e38, device=src.device, dtype=src.dtype) - index_expanded = index.unsqueeze(1).expand_as(src) - out.scatter_reduce_(0, index_expanded, src, reduce='amax', include_self=False) - return out - -batch_size = 1024 -features = 64 -dim_size = 128 - -def get_inputs(): - src = torch.randn(batch_size, features) - index = torch.randint(0, dim_size, (batch_size,)) - return [src, index] - -def get_init_inputs(): - return [dim_size] \ No newline at end of file