submission 102902

  from utils import make_match_reference, DeterministicContext
  import torch
  from task import input_t, output_t
+ from torch.utils.cpp_extension import load_inline
   
   
+ cpp_stub = r"""
+ #include <torch/extension.h>
+ void vector_add(torch::Tensor a, torch::Tensor b, torch::Tensor out);
+ """
+  
+ cuda_src = r"""
+ #include <torch/extension.h>
+ #include <cuda.h>
+ #include <cuda_runtime.h>
+ #include <cuda_fp16.h>
+ #include <ATen/cuda/CUDAContext.h>
+  
+ __global__ void vector_add_half2_kernel(const __half* __restrict__ a,
+                                         const __half* __restrict__ b,
+                                         __half* __restrict__ c,
+                                         int n) {
+     int idx2 = blockIdx.x * blockDim.x + threadIdx.x;  // half2 单位索引
+     int i = idx2 * 2;
+  
+     // 主路径：half2
+     if (i + 1 < n) {
+         const __half2* a2 = reinterpret_cast<const __half2*>(a);
+         const __half2* b2 = reinterpret_cast<const __half2*>(b);
+               __half2* c2 = reinterpret_cast<__half2*>(c);
+         __half2 ah = a2[idx2];
+         __half2 bh = b2[idx2];
+         c2[idx2] = __hadd2(ah, bh);
+     }
+  
+     // 尾元素（总元素为奇数时）
+     if ((n & 1) && idx2 == (n >> 1)) {
+         c[n - 1] = __hadd(a[n - 1], b[n - 1]);
+     }
+ }
+  
+ void vector_add(torch::Tensor a, torch::Tensor b, torch::Tensor out) {
+     TORCH_CHECK(a.is_cuda() && b.is_cuda() && out.is_cuda(), "tensors must be on CUDA");
+     TORCH_CHECK(a.dtype() == torch::kFloat16 && b.dtype() == torch::kFloat16 && out.dtype() == torch::kFloat16,
+                 "only float16 supported");
+     TORCH_CHECK(a.is_contiguous() && b.is_contiguous() && out.is_contiguous(),
+                 "tensors must be contiguous");
+     TORCH_CHECK(a.sizes() == b.sizes() && a.sizes() == out.sizes(), "shape mismatch");
+     TORCH_CHECK(a.dim() == 1 || a.dim() == 2, "expect 1D or 2D tensors");
+  
+     // 线性化处理（1D/2D 通用）
+     int n = a.numel();
+     int threads = 256;
+     int pairs = (n + 1) / 2;  // half2 元素数
+     int blocks = (pairs + threads - 1) / threads;
+  
+     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+     vector_add_half2_kernel<<<blocks, threads, 0, stream>>>(
+         reinterpret_cast<const __half*>(a.data_ptr<at::Half>()),
+         reinterpret_cast<const __half*>(b.data_ptr<at::Half>()),
+         reinterpret_cast<__half*>(out.data_ptr<at::Half>()),
+         n);
+     TORCH_CHECK(cudaGetLastError() == cudaSuccess, "kernel launch failed");
+ }
+ """
+  
+ major, minor = torch.cuda.get_device_capability()
+ # major, minor = 8, 9
+ arch_flag = f"-arch=sm_{major}{minor}"
+  
+ ext = load_inline(
+     name="vecadd_half2_2d_inline",
+     cpp_sources=[cpp_stub],
+     cuda_sources=[cuda_src],
+     functions=["vector_add"],
+     extra_cuda_cflags=[arch_flag],
+     verbose=True,
+ )
+  
+  
+ # def ref_kernel(data: input_t) -> output_t:
+ #     """
+ #     Reference implementation of vector addition using PyTorch.
+ #     Args:
+ #         data: Tuple of tensors [A, B] to be added.
+ #     Returns:
+ #         Tensor containing element-wise sums.
+ #     """
+ #     with DeterministicContext():
+ #         A, B, output = data
+ #         output[...] = A + B
+  
+ #         return output
+  
+  
  def custom_kernel(data: input_t) -> output_t:
      """
      Reference implementation of vector addition using PyTorch.
  ⋯ 4 unchanged lines
      """
      with DeterministicContext():
          A, B, output = data
-         # output[...] = A + B
-         torch.add(A, B, out=output)
+         ext.vector_add(A, B, output)
          return output
   
   
- def generate_input(size: int, seed: int) -> input_t:
-     """
-     Generates random input tensors of specified shapes.
-     Returns:
-         Tuple of tensors [A, B] 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()
-     C = torch.empty(size, size, device="cuda", dtype=torch.float16).contiguous()
-     return A, B, C
+ # def generate_input(size: int, seed: int) -> input_t:
+ #     """
+ #     Generates random input tensors of specified shapes.
+ #     Returns:
+ #         Tuple of tensors [A, B] 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()
+ #     C = torch.empty(size, size, device="cuda", dtype=torch.float16).contiguous()
+ #     return A, B, C
   
   
  # check_implementation = make_match_reference(ref_kernel)
   No newline at end of file