submission 67297

  from __future__ import annotations
   
- from typing import Tuple
+ from typing import Any, Dict, Tuple
   
  import torch
+ from torch.utils.cpp_extension import load_inline
   
  from task import input_t, output_t
   
   
+ CPP_DECL = r"""
+ #include <torch/extension.h>
+ void launch_vector_add(torch::Tensor A, torch::Tensor B, torch::Tensor C);
+ """
+  
+ CUDA_SRC = r"""
+ #include <ATen/cuda/CUDAContext.h>
+ #include <c10/cuda/CUDAGuard.h>
+ #include <c10/cuda/CUDAException.h>
+ #include <cuda_fp16.h>
+ #include <cstdint>
+  
+ namespace {
+  
+ constexpr int kBlockSize = 512;
+ constexpr int kMaxBlocksPerSM = 32;
+  
+ __global__ __launch_bounds__(kBlockSize, 2)
+ void vector_add_half2_kernel(const __half2* __restrict__ A,
+                              const __half2* __restrict__ B,
+                              __half2* __restrict__ C,
+                              int64_t pair_count,
+                              bool store_tail,
+                              const __half* __restrict__ A_scalar,
+                              const __half* __restrict__ B_scalar,
+                              __half* __restrict__ C_scalar,
+                              int64_t numel) {
+     const int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x;
+     int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+  
+     while (idx < pair_count) {
+         __half2 aval = __ldg(A + idx);
+         __half2 bval = __ldg(B + idx);
+         C[idx] = __hadd2(aval, bval);
+         idx += stride;
+     }
+  
+     if (store_tail && blockIdx.x == 0 && threadIdx.x == 0) {
+         const int64_t tail_index = numel - 1;
+         C_scalar[tail_index] = __hadd(A_scalar[tail_index], B_scalar[tail_index]);
+     }
+ }
+  
+ __global__ void vector_add_scalar_kernel(const __half* __restrict__ A,
+                                          const __half* __restrict__ B,
+                                          __half* __restrict__ C,
+                                          int64_t numel) {
+     const int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x;
+     int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+     while (idx < numel) {
+         C[idx] = __hadd(A[idx], B[idx]);
+         idx += stride;
+     }
+ }
+  
+ inline bool is_aligned_for_half2(const void* ptr) {
+     constexpr std::uintptr_t alignment = alignof(__half2);
+     return (reinterpret_cast<std::uintptr_t>(ptr) & (alignment - 1)) == 0;
+ }
+  
+ } // namespace
+  
+ void launch_vector_add(torch::Tensor A, torch::Tensor B, torch::Tensor C) {
+     TORCH_CHECK(A.is_cuda(), "A must be a CUDA tensor");
+     TORCH_CHECK(B.is_cuda(), "B must be a CUDA tensor");
+     TORCH_CHECK(C.is_cuda(), "C must be a CUDA tensor");
+  
+     TORCH_CHECK(A.scalar_type() == at::kHalf, "Expected torch.float16 tensors");
+     TORCH_CHECK(B.scalar_type() == at::kHalf, "Expected torch.float16 tensors");
+     TORCH_CHECK(C.scalar_type() == at::kHalf, "Expected torch.float16 tensors");
+  
+     TORCH_CHECK(A.is_contiguous(), "A must be contiguous");
+     TORCH_CHECK(B.is_contiguous(), "B must be contiguous");
+     TORCH_CHECK(C.is_contiguous(), "C must be contiguous");
+  
+     TORCH_CHECK(A.sizes() == B.sizes() && A.sizes() == C.sizes(),
+                 "Input and output tensors must have identical shapes");
+  
+     c10::cuda::CUDAGuard device_guard(A.device());
+     auto stream = at::cuda::getCurrentCUDAStream();
+  
+     const int64_t numel = A.numel();
+     if (numel == 0) {
+         return;
+     }
+  
+     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 bool use_half2 = numel > 1 &&
+         is_aligned_for_half2(A_ptr) &&
+         is_aligned_for_half2(B_ptr) &&
+         is_aligned_for_half2(C_ptr);
+  
+     auto* props = at::cuda::getCurrentDeviceProperties();
+     const int max_blocks = props->multiProcessorCount * kMaxBlocksPerSM;
+  
+     if (use_half2) {
+         const int64_t pair_count = numel >> 1;
+         const bool has_tail = (numel & 1) != 0;
+         int grid = static_cast<int>((pair_count + kBlockSize - 1) / kBlockSize);
+         if (grid < 1) {
+             grid = 1;
+         }
+         if (grid > max_blocks) {
+             grid = max_blocks;
+         }
+         vector_add_half2_kernel<<<grid, kBlockSize, 0, stream>>>(
+             reinterpret_cast<const __half2*>(A_ptr),
+             reinterpret_cast<const __half2*>(B_ptr),
+             reinterpret_cast<__half2*>(C_ptr),
+             pair_count,
+             has_tail,
+             A_ptr,
+             B_ptr,
+             C_ptr,
+             numel);
+     } else {
+         int grid = static_cast<int>((numel + kBlockSize - 1) / kBlockSize);
+         if (grid < 1) {
+             grid = 1;
+         }
+         if (grid > max_blocks) {
+             grid = max_blocks;
+         }
+         vector_add_scalar_kernel<<<grid, kBlockSize, 0, stream>>>(
+             A_ptr, B_ptr, C_ptr, numel);
+     }
+     C10_CUDA_KERNEL_LAUNCH_CHECK();
+ }
+ """
+  
+  
+ _kernel_cache: Dict[Tuple[int, int], Any] = {}
+  
+  
+ def _load_kernel(device: torch.device):
+     capability = torch.cuda.get_device_capability(device=device)
+     key = capability
+     if key in _kernel_cache:
+         return _kernel_cache[key]
+  
+     arch_flag = f"-gencode=arch=compute_{capability[0]}{capability[1]},code=sm_{capability[0]}{capability[1]}"
+     module = load_inline(
+         name=f"vector_add_cuda_{capability[0]}{capability[1]}",
+         cpp_sources=CPP_DECL,
+         cuda_sources=CUDA_SRC,
+         functions=["launch_vector_add"],
+         extra_cuda_cflags=[
+             "-O3",
+             "-std=c++17",
+             "-lineinfo",
+             "-use_fast_math",
+             arch_flag,
+         ],
+     )
+     _kernel_cache[key] = module
+     return module
+  
+  
  def _validate_tensors(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor) -> None:
      if not (A.is_cuda and B.is_cuda and C.is_cuda):
          raise RuntimeError("All tensors must reside on CUDA devices.")
  ⋯ 5 unchanged lines
          raise RuntimeError("Input and output tensors must be contiguous.")
   
   
- def _launch_vector_add(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor) -> None:
-     # `torch.add` issues an efficient pointwise kernel and honors the provided output buffer.
-     torch.add(A, B, out=C)
-  
-  
  def custom_kernel(data: input_t) -> output_t:
-     A, B, C = data  # type: ignore[misc]
+     try:
+         A, B, C = data  # type: ignore[misc]
+     except ValueError as err:
+         raise RuntimeError("Expected (A, B, output) tensor tuple.") from err
+  
      _validate_tensors(A, B, C)
-     _launch_vector_add(A, B, C)
+  
+     numel = A.numel()
+     if numel >= (1 << 20):
+         torch.add(A, B, out=C)
+         return C
+  
+     kernel = _load_kernel(A.device)
+     kernel.launch_vector_add(A, B, C)
      return C