submission 126189

- # !POPCORN leaderboard vectoradd_v2
- # !POPCORN gpus H100
- #
+ # #!POPCORN leaderboard vectoradd_v2
+ # #!POPCORN gpus L4
+  
  # #!/usr/bin/env python3
  # import torch
  # from torch.utils.cpp_extension import load_inline
- #
+  
  # # 1) C++ side: declaration so Python can call into the CUDA implementation
  # CPP_SRC = r"""
  # #include <torch/extension.h>
- #
+ # #include <cuda_runtime.h>
+  
  # // Forward declaration; definition will be in the CUDA translation unit.
  # torch::Tensor vecadd_cuda(torch::Tensor a, torch::Tensor b);
  # """
- #
- # # 2) CUDA side: kernel + definition (float32 math internally)
+  
+ # # 2) CUDA side: templated kernel + dtype dispatch
  # CUDA_SRC = r"""
  # #include <torch/extension.h>
  # #include <cuda.h>
  # #include <cuda_runtime.h>
- #
- # // 1D float32 vector add: c[i] = a[i] + b[i]
- # __global__ void vecadd_kernel(const float* __restrict__ a,
- #                               const float* __restrict__ b,
- #                               float* __restrict__ c,
+ # #include <ATen/ATen.h>
+  
+ # // 1D vector add: c[i] = a[i] + b[i] for arbitrary floating scalar_t
+ # template <typename scalar_t>
+ # __global__ void vecadd_kernel(const scalar_t* __restrict__ a,
+ #                               const scalar_t* __restrict__ b,
+ #                               scalar_t* __restrict__ c,
  #                               long n) {
- #     long idx = blockIdx.x * blockDim.x + threadIdx.x;
- #     if (idx < n) {
- #         c[idx] = a[idx] + b[idx];
+ #     long i = blockIdx.x * blockDim.x + threadIdx.x;
+ #     if (i < n) {
+ #         c[i] = a[i] + b[i];
  #     }
  # }
- #
+  
  # // Definition of vecadd_cuda declared in CPP_SRC
  # torch::Tensor vecadd_cuda(torch::Tensor a, torch::Tensor b) {
- #     auto a_c = a.contiguous();
- #     auto b_c = b.contiguous();
- #
+ #     // Assume inputs are already contiguous in this benchmark harness
+ #     auto a_c = a;
+ #     auto b_c = b;
+  
+ #     TORCH_CHECK(a_c.sizes() == b_c.sizes(), "Input tensors must have same shape");
+  
  #     auto out = torch::empty_like(a_c);
  #     long n = a_c.numel();
- #
- #     const float* a_ptr = a_c.data_ptr<float>();
- #     const float* b_ptr = b_c.data_ptr<float>();
- #     float* out_ptr     = out.data_ptr<float>();
- #
+  
+ #     if (n == 0) {
+ #         return out;
+ #     }
+  
  #     constexpr int threads = 256;
- #     int blocks = (int)((n + threads - 1) / threads);
- #
- #     vecadd_kernel<<<blocks, threads>>>(a_ptr, b_ptr, out_ptr, n);
- #
+ #     // Enough threads to cover all elements, 1 element per thread
+ #     int blocks = static_cast<int>((n + threads - 1) / threads);
+  
+ #     // Dispatch on the actual scalar type: float32, float16, bfloat16, double, ...
+ #     AT_DISPATCH_FLOATING_TYPES_AND_HALF(a_c.scalar_type(), "vecadd_cuda", [&] {
+ #         const scalar_t* a_ptr = a_c.data_ptr<scalar_t>();
+ #         const scalar_t* b_ptr = b_c.data_ptr<scalar_t>();
+ #         scalar_t* out_ptr     = out.data_ptr<scalar_t>();
+  
+ #         vecadd_kernel<scalar_t><<<blocks, threads>>>(
+ #             a_ptr, b_ptr, out_ptr, n
+ #         );
+ #     });
+  
  #     return out;
  # }
  # """
- #
+  
  # # 3) Build extension once at import time.
  # vecadd_mod = load_inline(
- #     name="vecadd_cuda_ext",
+ #     name="vecad",
  #     cpp_sources=CPP_SRC,
  #     cuda_sources=CUDA_SRC,
  #     functions=["vecadd_cuda"],
  #     with_cuda=True,
  #     verbose=False,
  # )
- #
- #
+  
+  
  # @torch.inference_mode()
  # def custom_kernel(data):
  #     """
  #     GPU Mode entrypoint.
- #
+  
  #     The harness passes a tuple with at least two elements, e.g.:
  #         (a, b, ...)
- #
+  
  #     We:
- #       - take the first two tensors,
- #       - upcast to float32 on CUDA,
+ #       - grab a and b,
+ #       - move them to CUDA (keeping their original dtype),
  #       - run the CUDA vecadd,
- #       - cast the result back to the original dtype of a,
- #       - and return a single tensor.
+ #       - return a single tensor.
+ #     No checks, no dtype conversions.
  #     """
- #     # 1) Extract a, b from whatever tuple/list the harness gives
- #     a = data[0]
- #     b = data[1]
- #
- #     orig_dtype = a.dtype
- #
- #     device = torch.device("cuda")
- #     a32 = a.to(device=device, dtype=torch.float32, non_blocking=True)
- #     b32 = b.to(device=device, dtype=torch.float32, non_blocking=True)
- #
- #
- #     # 5) Run CUDA vecadd on float32
- #     out32 = vecadd_mod.vecadd_cuda(a32, b32)
- #
- #     # 6) Cast back to original dtype to match the reference
- #     out = out32.to(dtype=orig_dtype)
- #
+ #     # a = data[0]  # keep this to prevent error
+ #     # b = data[1]  # keep this
+  
+ #     # Move to CUDA; keeps dtype, so we stay in orig_dtype
+ #     # a = a.cuda(non_blocking=True)
+ #     # b = b.cuda(non_blocking=True)
+  
+ #     out = vecadd_mod.vecadd_cuda(data[0], data[1])
+  
+ #     # No explicit synchronize: let the harness force sync when it needs the result
  #     return out
- #
   
  #!POPCORN leaderboard vectoradd_v2
  #!POPCORN gpus L4
   
  #!/usr/bin/env python3
  import torch
- from torch.utils.cpp_extension import load_inline
+ import triton
+ import triton.language as tl
   
- # 1) C++ side: declaration so Python can call into the CUDA implementation
- CPP_SRC = r"""
- #include <torch/extension.h>
   
- // Forward declaration; definition will be in the CUDA translation unit.
- torch::Tensor vecadd_cuda(torch::Tensor a, torch::Tensor b);
- """
+ @triton.jit
+ def vecadd_kernel(a_ptr, b_ptr, c_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+     """
+     Each Triton program (block) handles BLOCK_SIZE elements.
+     """
+     pid = tl.program_id(0)  # 1D grid
+     offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+     mask = offsets < n_elements
   
- # 2) CUDA side: templated kernel + dtype dispatch
- CUDA_SRC = r"""
- #include <torch/extension.h>
- #include <cuda.h>
- #include <cuda_runtime.h>
- #include <ATen/ATen.h>
+     a = tl.load(a_ptr + offsets, mask=mask)
+     b = tl.load(b_ptr + offsets, mask=mask)
+     c = a + b
+     tl.store(c_ptr + offsets, c, mask=mask)
   
- // 1D vector add: c[i] = a[i] + b[i] for arbitrary floating scalar_t
- template <typename scalar_t>
- __global__ void vecadd_kernel(const scalar_t* __restrict__ a,
-                               const scalar_t* __restrict__ b,
-                               scalar_t* __restrict__ c,
-                               long n) {
-     long idx = blockIdx.x * blockDim.x + threadIdx.x;
-     if (idx < n) {
-         c[idx] = a[idx] + b[idx];
-     }
- }
   
- // Definition of vecadd_cuda declared in CPP_SRC
- torch::Tensor vecadd_cuda(torch::Tensor a, torch::Tensor b) {
-     auto a_c = a.contiguous();
-     auto b_c = b.contiguous();
-  
-     auto out = torch::empty_like(a_c);
-     long n = a_c.numel();
-  
-     constexpr int threads = 256;
-     int blocks = (int)((n + threads - 1) / threads);
-  
-     // Dispatch on the actual scalar type: float32, float16, bfloat16, double, ...
-     AT_DISPATCH_FLOATING_TYPES_AND_HALF(a_c.scalar_type(), "vecadd_cuda", [&] {
-         const scalar_t* a_ptr = a_c.data_ptr<scalar_t>();
-         const scalar_t* b_ptr = b_c.data_ptr<scalar_t>();
-         scalar_t* out_ptr     = out.data_ptr<scalar_t>();
-  
-         vecadd_kernel<scalar_t><<<blocks, threads>>>(a_ptr, b_ptr, out_ptr, n);
-     });
-  
-     return out;
- }
- """
-  
- # 3) Build extension once at import time.
- vecadd_mod = load_inline(
-     name="vecadd_cuda_ext",
-     cpp_sources=CPP_SRC,
-     cuda_sources=CUDA_SRC,
-     functions=["vecadd_cuda"],
-     with_cuda=True,
-     verbose=False,
- )
-  
-  
  @torch.inference_mode()
  def custom_kernel(data):
      """
  ⋯ 5 unchanged lines
      We:
        - grab a and b,
        - move them to CUDA (keeping their original dtype),
-       - run the CUDA vecadd,
+       - run the Triton vecadd kernel,
        - return a single tensor.
-     No checks, no dtype conversions.
      """
-     a = data[0]
-     b = data[1]
+     a = data[0]  # keep this to prevent error
+     b = data[1]  # keep this
   
-     # Move to CUDA; keeps dtype, so we stay in orig_dtype
+     # Move to CUDA; keep original dtype
      a = a.cuda(non_blocking=True)
      b = b.cuda(non_blocking=True)
   
-     out = vecadd_mod.vecadd_cuda(a, b)
+     # Shapes must match
+     assert a.shape == b.shape, "Input tensors must have the same shape"
   
-     # No explicit synchronize: let the harness force sync when it needs the result
+     # Ensure contiguous (cheap if already contiguous)
+     a_c = a if a.is_contiguous() else a.contiguous()
+     b_c = b if b.is_contiguous() else b.contiguous()
+  
+     out = torch.empty_like(a_c)
+     n_elements = out.numel()
+  
+     if n_elements == 0:
+         return out
+  
+     BLOCK_SIZE = 1024  # elements per Triton program
+     grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
+  
+     # Launch Triton kernel
+     vecadd_kernel[grid](
+         a_c,
+         b_c,
+         out,
+         n_elements,
+         BLOCK_SIZE=BLOCK_SIZE,
+         num_warps=4,
+         num_stages=1,
+     )
+  
      return out