submission 67509

  from utils import make_match_reference, DeterministicContext
  import torch
+ from torch.utils.cpp_extension import load_inline
  from task import input_t, output_t
- import triton
- import triton.language as tl
   
+ vectoradd_source = r"""
+ #include <cuda_fp16.h>
+ #include <stdexcept>
   
- # Fixed optimal config - no autotuning variance
- @triton.jit
- def vecadd_fp16_kernel(A, B, C, N, BLOCK_SIZE: tl.constexpr):
-     pid = tl.program_id(0)
-     offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-     mask = offs < N
+ __global__ void __launch_bounds__(512, 2)
+ vectoradd_cuda_fast(
+     const half* __restrict__ A,
+     const half* __restrict__ B,
+     half* __restrict__ C,
+     const int N)
+ {
+     const int tid  = blockIdx.x * blockDim.x + threadIdx.x;
+     const int base = tid * 8;   // 8 elems per thread
   
-     # Load
-     a = tl.load(A + offs, mask=mask, other=0.0)
-     b = tl.load(B + offs, mask=mask, other=0.0)
+     if (base >= N) return;
   
-     # Compute
-     c = a + b
+     // fast path: we can read 8 halves (16B) safely
+     if (base + 7 < N) {
+         const uint4 a_vec = *reinterpret_cast<const uint4*>(A + base);
+         const uint4 b_vec = *reinterpret_cast<const uint4*>(B + base);
   
-     # Store
-     tl.store(C + offs, c, mask=mask)
+         const half2 a0 = reinterpret_cast<const half2&>(a_vec.x);
+         const half2 a1 = reinterpret_cast<const half2&>(a_vec.y);
+         const half2 a2 = reinterpret_cast<const half2&>(a_vec.z);
+         const half2 a3 = reinterpret_cast<const half2&>(a_vec.w);
   
+         const half2 b0 = reinterpret_cast<const half2&>(b_vec.x);
+         const half2 b1 = reinterpret_cast<const half2&>(b_vec.y);
+         const half2 b2 = reinterpret_cast<const half2&>(b_vec.z);
+         const half2 b3 = reinterpret_cast<const half2&>(b_vec.w);
   
- def triton_vecadd(A, B, C):
-     N = A.numel()
+         const half2 c0 = __hadd2(a0, b0);
+         const half2 c1 = __hadd2(a1, b1);
+         const half2 c2 = __hadd2(a2, b2);
+         const half2 c3 = __hadd2(a3, b3);
   
-     # Fixed optimal config based on your 235us result
-     # Tune BLOCK_SIZE based on what worked best in autotuning
-     BLOCK_SIZE = 4096  # Start with this, adjust based on your best run
-     num_warps = 16
+         uint4 c_vec;
+         reinterpret_cast<half2&>(c_vec.x) = c0;
+         reinterpret_cast<half2&>(c_vec.y) = c1;
+         reinterpret_cast<half2&>(c_vec.z) = c2;
+         reinterpret_cast<half2&>(c_vec.w) = c3;
   
-     grid = (triton.cdiv(N, BLOCK_SIZE),)
-     vecadd_fp16_kernel[grid](
-         A, B, C, N,
-         BLOCK_SIZE=BLOCK_SIZE,
-         num_warps=num_warps,
-         num_stages=1,
-     )
-     return C
+         *reinterpret_cast<uint4*>(C + base) = c_vec;
+     } else {
+         // tail: only final partial block hits this
+         #pragma unroll
+         for (int i = 0; i < 8; ++i) {
+             const int idx = base + i;
+             if (idx < N) {
+                 C[idx] = __hadd(A[idx], B[idx]);
+             }
+         }
+     }
+ }
   
+ // this is the function PyTorch calls
+ torch::Tensor vectoradd_triton_match(torch::Tensor A,
+                                      torch::Tensor B,
+                                      torch::Tensor C) {
+     const int N = A.numel();
   
+     const half* a_ptr = reinterpret_cast<const half*>(A.data_ptr<at::Half>());
+     const half* b_ptr = reinterpret_cast<const half*>(B.data_ptr<at::Half>());
+     half*       c_ptr = reinterpret_cast<half*>(C.data_ptr<at::Half>());
+  
+     const int threads = 512;
+     const int elems_per_block = 4096;
+     const int blocks = (N + elems_per_block - 1) / elems_per_block;
+  
+     // ✅ call the kernel we actually defined
+     vectoradd_cuda_fast<<<blocks, threads>>>(a_ptr, b_ptr, c_ptr, N);
+  
+     cudaError_t err = cudaGetLastError();
+     if (err != cudaSuccess) {
+         throw std::runtime_error(cudaGetErrorString(err));
+     }
+  
+     return C;
+ }
+ """
+  
+ vectoradd_cpp_source = r"""
+ #include <torch/extension.h>
+ torch::Tensor vectoradd_triton_match(torch::Tensor A,
+                                      torch::Tensor B,
+                                      torch::Tensor C);
+ """
+  
+ vectoradd_module = load_inline(
+     name='vectoradd_triton_match',
+     cpp_sources=vectoradd_cpp_source,
+     cuda_sources=vectoradd_source,
+     functions=['vectoradd_triton_match'],
+     verbose=False,
+     extra_cuda_cflags=[
+         '-O3',
+         '--use_fast_math',
+         '-gencode=arch=compute_100,code=sm_100',
+         '-Xptxas=-O3',
+     ],
+ )
+  
+  
  def ref_kernel(data: input_t) -> output_t:
-     """
-     Reference implementation of vector addition using PyTorch.
-     Args:
-         data: Tuple of tensors [A, B, output] to be added.
-     Returns:
-         Tensor containing element-wise sums.
-     """
      with DeterministicContext():
          A, B, output = data
          output[...] = A + B
  ⋯ 1 unchanged lines
   
   
  def generate_input(size: int, seed: int) -> input_t:
-     """
-     Generates random input tensors of specified shapes.
-     Returns:
-         Tuple of tensors [A, B, C] to be added.
-     """
      gen = torch.Generator(device="cuda")
      gen.manual_seed(seed)
-     A = torch.randn(
-         size, size, device="cuda", dtype=torch.float16, generator=gen
-     ).contiguous()
-     B = torch.randn(
-         size, size, device="cuda", dtype=torch.float16, generator=gen
-     ).contiguous()
+     A = torch.randn(size, size, device="cuda", dtype=torch.float16,
+                     generator=gen).contiguous()
+     B = torch.randn(size, size, device="cuda", dtype=torch.float16,
+                     generator=gen).contiguous()
      C = torch.empty(size, size, device="cuda", dtype=torch.float16).contiguous()
      return A, B, C
   
   
  def custom_kernel(data: input_t) -> output_t:
-     """Fixed optimal Triton config - no autotuning variance"""
      with DeterministicContext():
          A, B, C = data
-         return triton_vecadd(A, B, C)
+         return vectoradd_module.vectoradd_triton_match(A, B, C)
   
   
  check_implementation = make_match_reference(ref_kernel)