GPUCodeForces/S1/3/softmax_cuda.py

# softmax_cuda.py 
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

from softmax_torch import feature_dim

assert feature_dim % 4 == 0, "Feature dimension must be a multiple of 4 for float4 vectorization"

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 softmax_forward_cuda(torch::Tensor input);
        """

        cuda_source = f"""
        #include <cuda_runtime.h>
        #include <float.h>

        #define BLOCK_SIZE 512
        #define WARP_SIZE 32

        __device__ __forceinline__ float warp_reduce_max(float val) {{
            for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2)
                val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
            return val;
        }}

        __device__ __forceinline__ float warp_reduce_sum(float val) {{
            for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2)
                val += __shfl_down_sync(0xffffffff, val, offset);
            return val;
        }}

        // Optimized single-pass fused Softmax kernel with float4 vectorization
        __global__ void softmax_fused_vectorized_kernel(
            const float* __restrict__ input,
            float* __restrict__ output,
            int batch_size,
            int feature_dim
        ) {{
            extern __shared__ float sdata[];
            float* s_reducers = sdata;
            float* s_x_cache = &sdata[BLOCK_SIZE / WARP_SIZE];
            
            const int feature_dim_div4 = feature_dim / 4;

            int row = blockIdx.x;
            if (row >= batch_size) return;


            const float4* x4 = reinterpret_cast<const float4*>(input + row * feature_dim);
            float4* y4 = reinterpret_cast<float4*>(output + row * feature_dim);
            

            for (int i_vec = threadIdx.x; i_vec < feature_dim_div4; i_vec += BLOCK_SIZE) {{
                float4 val4 = x4[i_vec]; 
                float* cache_ptr = s_x_cache + i_vec * 4;

                cache_ptr[0] = val4.x;
                cache_ptr[1] = val4.y;
                cache_ptr[2] = val4.z;
                cache_ptr[3] = val4.w;
            }}
            __syncthreads();
            

            float thread_max = -FLT_MAX;
            for (int i = threadIdx.x; i < feature_dim; i += BLOCK_SIZE) {{
                thread_max = fmaxf(thread_max, s_x_cache[i]);
            }}

            float warp_max = warp_reduce_max(thread_max);
            int warp_id = threadIdx.x / WARP_SIZE;
            int lane_id = threadIdx.x % WARP_SIZE;
            if (lane_id == 0) s_reducers[warp_id] = warp_max;
            __syncthreads();

            thread_max = (threadIdx.x < BLOCK_SIZE / WARP_SIZE) ? s_reducers[lane_id] : -FLT_MAX;
            if (warp_id == 0) warp_max = warp_reduce_max(thread_max);
            
            if (threadIdx.x == 0) s_reducers[0] = warp_max;
            __syncthreads();
            float row_max = s_reducers[0];

            float thread_sum = 0.0f;
            for (int i = threadIdx.x; i < feature_dim; i += BLOCK_SIZE) {{
                thread_sum += expf(s_x_cache[i] - row_max);
            }}

            float warp_sum = warp_reduce_sum(thread_sum);
            if (lane_id == 0) s_reducers[warp_id] = warp_sum;
            __syncthreads();

            thread_sum = (threadIdx.x < BLOCK_SIZE / WARP_SIZE) ? s_reducers[lane_id] : 0.0f;
            if (warp_id == 0) warp_sum = warp_reduce_sum(thread_sum);

            if (threadIdx.x == 0) s_reducers[0] = warp_sum;
            __syncthreads();
            float row_sum = s_reducers[0];
            float inv_row_sum = 1.0f / row_sum; 

            for (int i_vec = threadIdx.x; i_vec < feature_dim_div4; i_vec += BLOCK_SIZE) {{
                float* cache_ptr = s_x_cache + i_vec * 4;
                
                float4 val4;
                
                // Read 4 scalars from shared memory and calculate
                val4.x = expf(cache_ptr[0] - row_max) * inv_row_sum;
                val4.y = expf(cache_ptr[1] - row_max) * inv_row_sum;
                val4.z = expf(cache_ptr[2] - row_max) * inv_row_sum;
                val4.w = expf(cache_ptr[3] - row_max) * inv_row_sum;

                y4[i_vec] = val4;
            }}
        }}

        torch::Tensor softmax_forward_cuda(torch::Tensor input) {{
            input = input.contiguous();
            int batch_size = input.size(0);
            int feature_dim = input.size(1);
            if (feature_dim % 4 != 0) {{
                AT_ERROR("Feature dimension must be a multiple of 4 for this kernel.");
            }}
            auto output = torch::empty_like(input);
            
            const int threads = BLOCK_SIZE;
            const int blocks = batch_size;
            
            size_t shared_mem_size = (BLOCK_SIZE / WARP_SIZE + feature_dim) * sizeof(float);
            
            softmax_fused_vectorized_kernel<<<blocks, threads, shared_mem_size>>>(
                input.data_ptr<float>(),
                output.data_ptr<float>(),
                batch_size,
                feature_dim
            );
            return output;
        }}
        """
        
        self.softmax_op = load_inline(
            name="softmax_fused_vectorized_op",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["softmax_forward_cuda"],
            extra_cuda_cflags=["-O3", "--use_fast_math"],
            verbose=False 
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.softmax_op.softmax_forward_cuda(x.contiguous())