submission 66403

  #!POPCORN leaderboard vectoradd_v2
  import torch
- import triton
- import triton.language as tl
+ from torch.utils.cpp_extension import load_inline
   
+ # CUDA kernel with vectorized memory access
+ cuda_source = """
+ #include <torch/extension.h>
+ #include <cuda_runtime.h>
   
- @triton.jit
- def add_kernel(
-     a_ptr,  # pointer to input A
-     b_ptr,  # pointer to input B
-     output_ptr,  # pointer to output
-     n_elements,  # total number of elements
-     BLOCK_SIZE: tl.constexpr,  # number of elements per block
- ):
-     """Optimized element-wise addition kernel using Triton"""
-     # Get program ID and compute element offset
-     pid = tl.program_id(axis=0)
-     block_start = pid * BLOCK_SIZE
-     offsets = block_start + tl.arange(0, BLOCK_SIZE)
+ // Vectorized kernel using float4 for 128-bit memory transactions
+ __global__ void vectorized_add_kernel(
+     const float* __restrict__ a,
+     const float* __restrict__ b,
+     float* __restrict__ output,
+     const int n
+ ) {
+     // Each thread processes 4 elements using float4
+     const int idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
      
-     # Create mask for boundary checking
-     mask = offsets < n_elements
+     if (idx + 3 < n) {
+         // Vectorized load (128-bit transaction)
+         float4 a_vec = *reinterpret_cast<const float4*>(a + idx);
+         float4 b_vec = *reinterpret_cast<const float4*>(b + idx);
+         
+         // Compute
+         float4 out_vec;
+         out_vec.x = a_vec.x + b_vec.x;
+         out_vec.y = a_vec.y + b_vec.y;
+         out_vec.z = a_vec.z + b_vec.z;
+         out_vec.w = a_vec.w + b_vec.w;
+         
+         // Vectorized store (128-bit transaction)
+         *reinterpret_cast<float4*>(output + idx) = out_vec;
+     }
+     else if (idx < n) {
+         // Handle remaining elements
+         for (int i = idx; i < n && i < idx + 4; i++) {
+             output[i] = a[i] + b[i];
+         }
+     }
+ }
+  
+ // Half precision (float16) vectorized kernel
+ __global__ void vectorized_add_kernel_half(
+     const __half* __restrict__ a,
+     const __half* __restrict__ b,
+     __half* __restrict__ output,
+     const int n
+ ) {
+     // Each thread processes 8 elements using float4 (which holds 8 halfs)
+     const int idx = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
      
-     # Load data with masking
-     a = tl.load(a_ptr + offsets, mask=mask)
-     b = tl.load(b_ptr + offsets, mask=mask)
+     if (idx + 7 < n) {
+         // Load 8 half values as float4 (128-bit)
+         float4 a_vec = *reinterpret_cast<const float4*>(a + idx);
+         float4 b_vec = *reinterpret_cast<const float4*>(b + idx);
+         
+         // Reinterpret as half2 for computation
+         __half2* a_h2 = reinterpret_cast<__half2*>(&a_vec);
+         __half2* b_h2 = reinterpret_cast<__half2*>(&b_vec);
+         __half2 out_h2[4];
+         
+         #pragma unroll
+         for (int i = 0; i < 4; i++) {
+             out_h2[i] = __hadd2(a_h2[i], b_h2[i]);
+         }
+         
+         // Store result
+         *reinterpret_cast<float4*>(output + idx) = *reinterpret_cast<float4*>(out_h2);
+     }
+     else if (idx < n) {
+         // Handle remaining elements
+         for (int i = idx; i < n && i < idx + 8; i++) {
+             output[i] = __hadd(a[i], b[i]);
+         }
+     }
+ }
+  
+ torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor output) {
+     const int n = a.numel();
      
-     # Perform addition
-     output = a + b
+     if (a.dtype() == torch::kFloat32) {
+         // Use float4 vectorization (4 floats per thread)
+         const int threads = 256;
+         const int blocks = (n + threads * 4 - 1) / (threads * 4);
+         
+         vectorized_add_kernel<<<blocks, threads>>>(
+             a.data_ptr<float>(),
+             b.data_ptr<float>(),
+             output.data_ptr<float>(),
+             n
+         );
+     }
+     else if (a.dtype() == torch::kFloat16) {
+         // Use float4 vectorization (8 halfs per thread)
+         const int threads = 256;
+         const int blocks = (n + threads * 8 - 1) / (threads * 8);
+         
+         vectorized_add_kernel_half<<<blocks, threads>>>(
+             reinterpret_cast<const __half*>(a.data_ptr<at::Half>()),
+             reinterpret_cast<const __half*>(b.data_ptr<at::Half>()),
+             reinterpret_cast<__half*>(output.data_ptr<at::Half>()),
+             n
+         );
+     }
      
-     # Store result with masking
-     tl.store(output_ptr + offsets, output, mask=mask)
+     return output;
+ }
+ """
   
+ cpp_source = """
+ torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor output);
+ """
   
+ # Compile the CUDA extension
+ try:
+     cuda_module = load_inline(
+         name='vectorized_add',
+         cpp_sources=cpp_source,
+         cuda_sources=cuda_source,
+         functions=['add_cuda'],
+         extra_cuda_cflags=['-O3', '--use_fast_math', '-lineinfo'],
+         verbose=False,
+     )
+ except Exception as e:
+     print(f"CUDA compilation failed: {e}")
+     cuda_module = None
+  
+  
  def custom_kernel(data):
      """
-     High-performance vector addition using Triton CUDA kernel.
+     High-performance vector addition using raw CUDA with vectorized memory access.
+     
+     Optimizations:
+     - float4 vectorized loads/stores (128-bit memory transactions)
+     - __half2 vectorized arithmetic for fp16
+     - Fast math optimizations
+     - Coalesced memory access pattern
+     
      Args:
          data: Tuple of tensors (A, B, output)
      Returns:
  ⋯ 1 unchanged lines
      """
      A, B, output = data
      
-     # Get total number of elements
-     n_elements = output.numel()
-     
-     # Choose optimal block size (tuned for modern GPUs)
-     BLOCK_SIZE = 1024
-     
-     # Calculate grid size
-     grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
-     
-     # Launch kernel
-     add_kernel[grid](
-         A, B, output,
-         n_elements,
-         BLOCK_SIZE=BLOCK_SIZE,
-     )
-     
-     return output
   No newline at end of file
+     if cuda_module is not None:
+         return cuda_module.add_cuda(A, B, output)
+     else:
+         # Fallback to PyTorch if CUDA compilation failed
+         output[...] = A + B
+         return output
   No newline at end of file