submission 68288

- from utils import make_match_reference, DeterministicContext
- import torch
- from torch.utils.cpp_extension import load_inline
- from task import input_t, output_t
- 
- # CUDA: no-tail, assumes N % 4096 == 0
- vectoradd_source = r"""
- #include <cuda_fp16.h>
- #include <stdexcept>
- 
- // 512 threads, 8 fp16 elements per thread = 4096 elements per block
- // we assume N is divisible by 4096, so no tail / no bounds
- __global__ void __launch_bounds__(1024, 2)
- vectoradd_cuda_fast_nt(
-     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 halves = 16 bytes
- 
-     // no bounds checks — N is multiple of 4096
- 
-     // 16B load from A and B
-     const uint4 a = *reinterpret_cast<const uint4*>(A + base);
-     const uint4 b = *reinterpret_cast<const uint4*>(B + base);
- 
-     // unpack as half2
-     const half2 a0 = reinterpret_cast<const half2&>(a.x);
-     const half2 a1 = reinterpret_cast<const half2&>(a.y);
-     const half2 a2 = reinterpret_cast<const half2&>(a.z);
-     const half2 a3 = reinterpret_cast<const half2&>(a.w);
- 
-     const half2 b0 = reinterpret_cast<const half2&>(b.x);
-     const half2 b1 = reinterpret_cast<const half2&>(b.y);
-     const half2 b2 = reinterpret_cast<const half2&>(b.z);
-     const half2 b3 = reinterpret_cast<const half2&>(b.w);
- 
-     // add
-     uint4 c;
-     reinterpret_cast<half2&>(c.x) = __hadd2(a0, b0);
-     reinterpret_cast<half2&>(c.y) = __hadd2(a1, b1);
-     reinterpret_cast<half2&>(c.z) = __hadd2(a2, b2);
-     reinterpret_cast<half2&>(c.w) = __hadd2(a3, b3);
- 
-     // store
-     *reinterpret_cast<uint4*>(C + base) = c;
- }
- 
- // C++ binding that 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;  // 512 * 8
-     const int blocks = (N + elems_per_block - 1) / elems_per_block;
- 
-     // straight launch, no error check
-     vectoradd_cuda_fast_nt<<<blocks, threads>>>(a_ptr, b_ptr, c_ptr, N);
-     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',
-         # H100 / Hopper
-         '-gencode=arch=compute_90,code=sm_90',
-         # B200 / Blackwell
-         '-gencode=arch=compute_100,code=sm_100',
-         # stream through L2, don't clutter L1
-         '-Xptxas=-O3,-dlcm=cg',
-     ],
- )
- 
- 
- def ref_kernel(data: input_t) -> output_t:
-     # pure PyTorch reference
-     with DeterministicContext():
-         A, B, output = data
-         output[...] = A + B
-         return output
- 
- 
- def generate_input(size: int, seed: int) -> input_t:
-     # assuming square, e.g. 16384
-     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()
-     C = torch.empty(size, size, device="cuda", dtype=torch.float16).contiguous()
-     return A, B, C
- 
- 
- def custom_kernel(data: input_t) -> output_t:
-     with DeterministicContext():
-         A, B, C = data
-         return vectoradd_module.vectoradd_triton_match(A, B, C)
- 
- 
- check_implementation = make_match_reference(ref_kernel)
+ from utils import make_match_reference, DeterministicContext
+ import torch
+ from torch.utils.cpp_extension import load_inline
+ from task import input_t, output_t
+  
+ vectoradd_source = r"""
+ #include <cuda_fp16.h>
+ #include <stdexcept>
+  
+ __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
+  
+     if (base >= N) return;
+  
+     // 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);
+  
+         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);
+  
+         const half2 c0 = __hadd2(a0, b0);
+         const half2 c1 = __hadd2(a1, b1);
+         const half2 c2 = __hadd2(a2, b2);
+         const half2 c3 = __hadd2(a3, b3);
+  
+         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;
+  
+         *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:
+     with DeterministicContext():
+         A, B, output = data
+         output[...] = A + B
+         return output
+  
+  
+ def generate_input(size: int, seed: int) -> input_t:
+     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()
+     C = torch.empty(size, size, device="cuda", dtype=torch.float16).contiguous()
+     return A, B, C
+  
+  
+ def custom_kernel(data: input_t) -> output_t:
+     with DeterministicContext():
+         A, B, C = data
+         return vectoradd_module.vectoradd_triton_match(A, B, C)
+  
+  
+ check_implementation = make_match_reference(ref_kernel)