submission 67548

+ #generally solid triton script.
  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>
   
- __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
+ # 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
   
-     if (base >= N) return;
+     # Load
+     a = tl.load(A + offs, mask=mask, other=0.0)
+     b = tl.load(B + offs, mask=mask, other=0.0)
   
-     // 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);
+     # Compute
+     c = a + b
   
-         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);
+     # Store
+     tl.store(C + offs, c, mask=mask)
   
-         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);
+ def triton_vecadd(A, B, C):
+     N = A.numel()
   
-         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;
+     # 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
   
-         *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]);
-             }
-         }
-     }
- }
+     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
   
- // 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 vectoradd_module.vectoradd_triton_match(A, B, C)
+         return triton_vecadd(A, B, C)
   
   
  check_implementation = make_match_reference(ref_kernel)