submission 66563

  import torch
- import triton
- import triton.language as tl
+ from torch.utils.cpp_extension import load_inline
   
- from task import input_t, output_t
+ cuda_src = r"""
+ #include <torch/extension.h>
+ #include <ATen/cuda/CUDAContext.h>
+ #include <c10/cuda/CUDAGuard.h>
   
- DEVICE = triton.runtime.driver.active.get_active_torch_device()
+ #include <cuda.h>
+ #include <cuda_runtime.h>
+ #include <cuda_fp16.h>
+ #include <stdint.h>
   
- # ---- Autotuned, vectorized kernel ------------------------------------------
- @triton.autotune(
-     configs=[
-         # Good starting points across Ampere/Ada/Hopper; feel free to extend.
-         triton.Config({'BLOCK_SIZE': 256,  'VEC': 8}, num_warps=4,  num_stages=2),
-         triton.Config({'BLOCK_SIZE': 512,  'VEC': 8}, num_warps=8,  num_stages=2),
-         triton.Config({'BLOCK_SIZE': 1024, 'VEC': 4}, num_warps=8,  num_stages=2),
-         triton.Config({'BLOCK_SIZE': 2048, 'VEC': 2}, num_warps=8,  num_stages=2),
-         triton.Config({'BLOCK_SIZE': 4096, 'VEC': 1}, num_warps=8,  num_stages=2),
-     ],
-     key=['n_elements'],  # autotune choice depends mainly on problem size
- )
- @triton.jit
- def add_kernel(x_ptr, y_ptr, out_ptr, n_elements,
-                BLOCK_SIZE: tl.constexpr, VEC: tl.constexpr):
-     pid = tl.program_id(axis=0)
+ // Helpers --------------------------------------------------------------------
+ static inline int div_up_long(long long n, int d) {
+     return (int)((n + d - 1) / d);
+ }
   
-     # vectorized block start; each program handles BLOCK_SIZE * VEC elements
-     block_start = pid * BLOCK_SIZE * VEC
+ // Union to safely bit-cast between u32 and __half2
+ union u32_h2 {
+     uint32_t u;
+     __half2  h2;
+ };
   
-     # 2D tile: [BLOCK_SIZE, VEC] so compiler emits wider loads/stores
-     offs = block_start + tl.arange(0, BLOCK_SIZE)[:, None] * VEC + tl.arange(0, VEC)[None, :]
-     mask = offs < n_elements
+ // fast path: 128-bit vectorized (8 x fp16 per pack), requires 16B alignment
+ __global__ void add_fp16_vec8_kernel(const __half* __restrict__ a,
+                                      const __half* __restrict__ b,
+                                      __half* __restrict__ out,
+                                      size_t n_packs /* number of 8-half packs */) {
+     const uint4* __restrict__ a4 = reinterpret_cast<const uint4*>(a);
+     const uint4* __restrict__ b4 = reinterpret_cast<const uint4*>(b);
+     uint4* __restrict__ out4     = reinterpret_cast<uint4*>(out);
   
-     # Optional: tell compiler about alignment/contiguity (safe for contiguous tensors)
-     tl.multiple_of(block_start, VEC)
+     const size_t tid    = blockIdx.x * blockDim.x + threadIdx.x;
+     const size_t stride = (size_t)blockDim.x * gridDim.x;
   
-     x = tl.load(x_ptr + offs, mask=mask, other=0.0)
-     y = tl.load(y_ptr + offs, mask=mask, other=0.0)
-     z = x + y
-     tl.store(out_ptr + offs, z, mask=mask)
+     // ILP to increase memory-level parallelism per thread
+     const int ILP = 4;
+     size_t i = tid;
   
+     for (; i + (size_t)(ILP - 1) * stride < n_packs; i += (size_t)ILP * stride) {
+         #pragma unroll
+         for (int k = 0; k < ILP; ++k) {
+             size_t idx = i + (size_t)k * stride;
+             uint4 av = a4[idx];
+             uint4 bv = b4[idx];
   
- def add_into(x: torch.Tensor, y: torch.Tensor, out: torch.Tensor):
-     # Preconditions
-     assert x.device == DEVICE and y.device == DEVICE and out.device == DEVICE
-     assert x.is_contiguous() and y.is_contiguous() and out.is_contiguous()
-     assert x.shape == y.shape == out.shape
-     # (Your task already restricts to fp16, which is great for bandwidth.)
-     assert x.dtype == torch.float16 and y.dtype == torch.float16 and out.dtype == torch.float16
+             // Each uint4.{x,y,z,w} is 32-bit = __half2
+             u32_h2 a0; a0.u = av.x; u32_h2 b0; b0.u = bv.x; u32_h2 r0; r0.h2 = __hadd2(a0.h2, b0.h2);
+             u32_h2 a1; a1.u = av.y; u32_h2 b1; b1.u = bv.y; u32_h2 r1; r1.h2 = __hadd2(a1.h2, b1.h2);
+             u32_h2 a2; a2.u = av.z; u32_h2 b2; b2.u = bv.z; u32_h2 r2; r2.h2 = __hadd2(a2.h2, b2.h2);
+             u32_h2 a3; a3.u = av.w; u32_h2 b3; b3.u = bv.w; u32_h2 r3; r3.h2 = __hadd2(a3.h2, b3.h2);
   
-     n_elements = out.numel()
-     # Each program covers BLOCK_SIZE * VEC elements (meta values chosen by autotune).
-     grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE'] * meta['VEC']),)
-     add_kernel[grid](x, y, out, n_elements)  # meta-params picked by autotune
+             uint4 rv;
+             rv.x = r0.u; rv.y = r1.u; rv.z = r2.u; rv.w = r3.u;
+             out4[idx] = rv;
+         }
+     }
+     for (; i < n_packs; i += stride) {
+         uint4 av = a4[i];
+         uint4 bv = b4[i];
+         u32_h2 a0; a0.u = av.x; u32_h2 b0; b0.u = bv.x; u32_h2 r0; r0.h2 = __hadd2(a0.h2, b0.h2);
+         u32_h2 a1; a1.u = av.y; u32_h2 b1; b1.u = bv.y; u32_h2 r1; r1.h2 = __hadd2(a1.h2, b1.h2);
+         u32_h2 a2; a2.u = av.z; u32_h2 b2; b2.u = bv.z; u32_h2 r2; r2.h2 = __hadd2(a2.h2, b2.h2);
+         u32_h2 a3; a3.u = av.w; u32_h2 b3; b3.u = bv.w; u32_h2 r3; r3.h2 = __hadd2(a3.h2, b3.h2);
+         uint4 rv; rv.x = r0.u; rv.y = r1.u; rv.z = r2.u; rv.w = r3.u;
+         out4[i] = rv;
+     }
+ }
   
+ // medium path: __half2 (4B), requires 4B alignment; 2 x fp16 per element
+ __global__ void add_fp16_h2_kernel(const __half2* __restrict__ a,
+                                    const __half2* __restrict__ b,
+                                    __half2* __restrict__ out,
+                                    size_t n_h2) {
+     const size_t tid    = blockIdx.x * blockDim.x + threadIdx.x;
+     const size_t stride = (size_t)blockDim.x * gridDim.x;
+     for (size_t i = tid; i < n_h2; i += stride) {
+         out[i] = __hadd2(a[i], b[i]);
+     }
+ }
   
- def custom_kernel(data: input_t) -> output_t:
-     """
-     Custom implementation of vector addition using Triton (optimized).
-     Args:
-         inputs: triple (A, B, output)
-     Returns:
-         output (element-wise sum)
-     """
+ // tail / worst-case path: scalar half (handles misalignment + remainder)
+ __global__ void add_fp16_scalar_kernel(const __half* __restrict__ a,
+                                        const __half* __restrict__ b,
+                                        __half* __restrict__ out,
+                                        size_t n) {
+     const size_t tid    = blockIdx.x * blockDim.x + threadIdx.x;
+     const size_t stride = (size_t)blockDim.x * gridDim.x;
+     for (size_t i = tid; i < n; i += stride) {
+         out[i] = __hadd(a[i], b[i]);
+     }
+ }
+  
+ void add_fp16_out(torch::Tensor a, torch::Tensor b, torch::Tensor out) {
+     TORCH_CHECK(a.is_cuda() && b.is_cuda() && out.is_cuda(), "Tensors must be CUDA");
+     TORCH_CHECK(a.scalar_type() == torch::kHalf && b.scalar_type() == torch::kHalf && out.scalar_type() == torch::kHalf,
+                 "Tensors must be torch.float16");
+     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(), "Shapes must match");
+  
+     c10::cuda::CUDAGuard device_guard(a.get_device());
+     auto* a_ptr = reinterpret_cast<const __half*>(a.data_ptr<at::Half>());
+     auto* b_ptr = reinterpret_cast<const __half*>(b.data_ptr<at::Half>());
+     auto* o_ptr = reinterpret_cast<__half*>(out.data_ptr<at::Half>());
+  
+     const size_t n = (size_t)out.numel();
+     if (n == 0) return;
+  
+     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  
+     // Choose launch geometry
+     const int threads = 256;
+     const int sm = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+     int blocks;
+  
+     // Check alignments
+     uintptr_t a_addr = reinterpret_cast<uintptr_t>(a_ptr);
+     uintptr_t b_addr = reinterpret_cast<uintptr_t>(b_ptr);
+     uintptr_t o_addr = reinterpret_cast<uintptr_t>(o_ptr);
+  
+     const bool aligned16 = ((a_addr | b_addr | o_addr) & 0xF) == 0; // 16B
+     const bool aligned4  = ((a_addr | b_addr | o_addr) & 0x3) == 0; // 4B
+  
+     if (aligned16) {
+         const size_t n_packs = n / 8;  // 8 half per 128-bit pack
+         if (n_packs) {
+             blocks = std::min(div_up_long((long long)n_packs, threads), sm * 32);
+             add_fp16_vec8_kernel<<<blocks, threads, 0, stream>>>(a_ptr, b_ptr, o_ptr, n_packs);
+         }
+         const size_t tail = n - n_packs * 8;
+         if (tail) {
+             const __half* a_tail = a_ptr + n_packs * 8;
+             const __half* b_tail = b_ptr + n_packs * 8;
+             __half*       o_tail = o_ptr + n_packs * 8;
+             blocks = std::min(div_up_long((long long)tail, threads), sm * 4);
+             add_fp16_scalar_kernel<<<blocks, threads, 0, stream>>>(a_tail, b_tail, o_tail, tail);
+         }
+     } else if (aligned4) {
+         const size_t n_h2 = n / 2;
+         if (n_h2) {
+             blocks = std::min(div_up_long((long long)n_h2, threads), sm * 32);
+             const __half2* a2 = reinterpret_cast<const __half2*>(a_ptr);
+             const __half2* b2 = reinterpret_cast<const __half2*>(b_ptr);
+             __half2*       o2 = reinterpret_cast<__half2*>(o_ptr);
+             add_fp16_h2_kernel<<<blocks, threads, 0, stream>>>(a2, b2, o2, n_h2);
+         }
+         if (n & 1) {
+             const __half* a_tail = a_ptr + (n_h2 * 2);
+             const __half* b_tail = b_ptr + (n_h2 * 2);
+             __half*       o_tail = o_ptr + (n_h2 * 2);
+             add_fp16_scalar_kernel<<<1, 32, 0, stream>>>(a_tail, b_tail, o_tail, 1);
+         }
+     } else {
+         blocks = std::min(div_up_long((long long)n, threads), sm * 32);
+         add_fp16_scalar_kernel<<<blocks, threads, 0, stream>>>(a_ptr, b_ptr, o_ptr, n);
+     }
+  
+     TORCH_CHECK(cudaGetLastError() == cudaSuccess, "CUDA kernel launch failed");
+ }
+  
+ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+     m.def("add_fp16_out", &add_fp16_out, "Vector add (fp16) into out (CUDA)");
+ }
+ """
+  
+ ext = load_inline(
+     name="vecadd_fp16_inline",
+     cpp_sources="",
+     cuda_sources=cuda_src,
+     extra_cflags=["-O3"],
+     extra_cuda_cflags=[
+         "-O3",
+         "--use_fast_math",
+         # Enable half/half2 intrinsics explicitly:
+         "-U__CUDA_NO_HALF_OPERATORS__",
+         "-U__CUDA_NO_HALF_CONVERSIONS__",
+         "-U__CUDA_NO_HALF2_OPERATORS__",   # <-- added
+         # "-maxrregcount=64",  # optional tuning
+     ],
+     verbose=False,
+     with_cuda=True,
+ )
+  
+ def add_cuda_(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor):
+     assert A.is_cuda and B.is_cuda and out.is_cuda
+     assert A.dtype == torch.float16 and B.dtype == torch.float16 and out.dtype == torch.float16
+     assert A.shape == B.shape == out.shape
+     assert A.is_contiguous() and B.is_contiguous() and out.is_contiguous()
+     ext.add_fp16_out(A, B, out)
+     return out
+  
+ def add_cuda(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+     out = torch.empty_like(A)
+     return add_cuda_(A, B, out)
+  
+  
+ def custom_kernel(data):
      A, B, output = data
-     assert A.is_cuda and B.is_cuda and output.is_cuda, "Tensors must be on GPU"
-     assert A.shape == B.shape == output.shape, "All shapes must match"
-     assert A.dtype == torch.float16 and B.dtype == torch.float16 and output.dtype == torch.float16, \
-         "All tensors must be float16"
-     add_into(A, B, output)
+     add_cuda_(A, B, output)   # writes directly into 'output'
      return output