submission 67853

  ⋯ 1 unchanged lines
  from torch.utils.cpp_extension import load_inline
  from task import input_t, output_t
   
- cuda_source = """
- #include <cuda_fp16.h>
+ _cuda_src = r'''
  #include <torch/extension.h>
+ #include <ATen/cuda/CUDAContext.h>
+ #include <c10/cuda/CUDAGuard.h>
+ #include <cuda_fp16.h>
+ #include <stdint.h>
   
- __global__ void __launch_bounds__(256, 4) vectoradd_kernel(
-     const half* __restrict__ a,
-     const half* __restrict__ b,
-     half* __restrict__ out,
-     const int n)
- {
-     // Coalesced: consecutive threads access consecutive memory
-     int idx = blockIdx.x * blockDim.x + threadIdx.x;
-     int base = idx * 16;  // 16 elements per thread for coalescing
+ template<int VEC>
+ __global__ void vec_add_half2_ptx(const __half* __restrict__ A,
+                                   const __half* __restrict__ B,
+                                   __half* __restrict__ C,
+                                   int64_t N) {
+     const int64_t idx0 = (int64_t(blockIdx.x) * blockDim.x + threadIdx.x) * VEC;
+     const int64_t stride = int64_t(blockDim.x) * gridDim.x * VEC;
   
-     if (base + 15 < n) {
-         uint4 a0, a1, b0, b1, c0, c1;
+     const uintptr_t baseA = reinterpret_cast<uintptr_t>(A);
+     const uintptr_t baseB = reinterpret_cast<uintptr_t>(B);
+     const uintptr_t baseC = reinterpret_cast<uintptr_t>(C);
+     const bool aligned16 = (((baseA | baseB | baseC) & 0xF) == 0);
   
-         const half* ap = a + base;
-         const half* bp = b + base;
-         half* cp = out + base;
+     for (int64_t idx = idx0; idx < N; idx += stride) {
+         int64_t remaining = N - idx;
+         if (remaining >= VEC) {
+             const char* a_ptr = reinterpret_cast<const char*>(A) + idx * 2;
+             const char* b_ptr = reinterpret_cast<const char*>(B) + idx * 2;
+             char*       c_ptr = reinterpret_cast<char*>(C) + idx * 2;
   
-         // Load 16 fp16 (2x 128-bit)
-         asm volatile ("ld.global.ca.v4.u32 {%0,%1,%2,%3},[%4];" : "=r"(a0.x),"=r"(a0.y),"=r"(a0.z),"=r"(a0.w) : "l"(ap));
-         asm volatile ("ld.global.ca.v4.u32 {%0,%1,%2,%3},[%4];" : "=r"(a1.x),"=r"(a1.y),"=r"(a1.z),"=r"(a1.w) : "l"(ap+8));
-         asm volatile ("ld.global.ca.v4.u32 {%0,%1,%2,%3},[%4];" : "=r"(b0.x),"=r"(b0.y),"=r"(b0.z),"=r"(b0.w) : "l"(bp));
-         asm volatile ("ld.global.ca.v4.u32 {%0,%1,%2,%3},[%4];" : "=r"(b1.x),"=r"(b1.y),"=r"(b1.z),"=r"(b1.w) : "l"(bp+8));
+             if (aligned16) {
+                 // 128-bit vectorized path: load/store 16 bytes per thread, do 4x add.f16x2
+                 unsigned rA0, rA1, rA2, rA3;
+                 unsigned rB0, rB1, rB2, rB3;
+                 unsigned rC0, rC1, rC2, rC3;
+                 asm volatile(
+                     "{\n\t"
+                     "ld.global.v4.b32 {%0,%1,%2,%3}, [%12];\n\t"
+                     "ld.global.v4.b32 {%4,%5,%6,%7}, [%13];\n\t"
+                     "add.f16x2 %8,  %0, %4;\n\t"
+                     "add.f16x2 %9,  %1, %5;\n\t"
+                     "add.f16x2 %10, %2, %6;\n\t"
+                     "add.f16x2 %11, %3, %7;\n\t"
+                     "st.global.v4.b32 [%14], {%8,%9,%10,%11};\n\t"
+                     "}\n"
+                     : "=&r"(rA0), "=&r"(rA1), "=&r"(rA2), "=&r"(rA3),
+                       "=&r"(rB0), "=&r"(rB1), "=&r"(rB2), "=&r"(rB3),
+                       "=&r"(rC0), "=&r"(rC1), "=&r"(rC2), "=&r"(rC3)
+                     : "l"(a_ptr), "l"(b_ptr), "l"(c_ptr)
+                     : "memory");
+             } else {
+                 // 32-bit path: tolerate only 4B alignment using scalar b32 ops, still using add.f16x2
+                 unsigned rA0, rA1, rA2, rA3;
+                 unsigned rB0, rB1, rB2, rB3;
+                 unsigned rC0, rC1, rC2, rC3;
   
-         // SIMD add
-         asm volatile ("add.f16x2 %0,%1,%2;" : "=r"(c0.x) : "r"(a0.x), "r"(b0.x));
-         asm volatile ("add.f16x2 %0,%1,%2;" : "=r"(c0.y) : "r"(a0.y), "r"(b0.y));
-         asm volatile ("add.f16x2 %0,%1,%2;" : "=r"(c0.z) : "r"(a0.z), "r"(b0.z));
-         asm volatile ("add.f16x2 %0,%1,%2;" : "=r"(c0.w) : "r"(a0.w), "r"(b0.w));
-         asm volatile ("add.f16x2 %0,%1,%2;" : "=r"(c1.x) : "r"(a1.x), "r"(b1.x));
-         asm volatile ("add.f16x2 %0,%1,%2;" : "=r"(c1.y) : "r"(a1.y), "r"(b1.y));
-         asm volatile ("add.f16x2 %0,%1,%2;" : "=r"(c1.z) : "r"(a1.z), "r"(b1.z));
-         asm volatile ("add.f16x2 %0,%1,%2;" : "=r"(c1.w) : "r"(a1.w), "r"(b1.w));
+                 const char* a0 = a_ptr + 0;
+                 const char* a1 = a_ptr + 4;
+                 const char* a2 = a_ptr + 8;
+                 const char* a3 = a_ptr + 12;
   
-         // Store
-         asm volatile ("st.global.cg.v4.u32 [%0],{%1,%2,%3,%4};" :: "l"(cp),"r"(c0.x),"r"(c0.y),"r"(c0.z),"r"(c0.w));
-         asm volatile ("st.global.cg.v4.u32 [%0],{%1,%2,%3,%4};" :: "l"(cp+8),"r"(c1.x),"r"(c1.y),"r"(c1.z),"r"(c1.w));
-     }
-     else if (base < n) {
-         for (int i = 0; i < 16 && base + i < n; i++) {
-             out[base + i] = __hadd(a[base + i], b[base + i]);
+                 const char* b0 = b_ptr + 0;
+                 const char* b1 = b_ptr + 4;
+                 const char* b2 = b_ptr + 8;
+                 const char* b3 = b_ptr + 12;
+  
+                 char* c0 = c_ptr + 0;
+                 char* c1 = c_ptr + 4;
+                 char* c2 = c_ptr + 8;
+                 char* c3 = c_ptr + 12;
+  
+                 asm volatile("ld.global.b32 %0, [%1];" : "=&r"(rA0) : "l"(a0) : "memory");
+                 asm volatile("ld.global.b32 %0, [%1];" : "=&r"(rA1) : "l"(a1) : "memory");
+                 asm volatile("ld.global.b32 %0, [%1];" : "=&r"(rA2) : "l"(a2) : "memory");
+                 asm volatile("ld.global.b32 %0, [%1];" : "=&r"(rA3) : "l"(a3) : "memory");
+  
+                 asm volatile("ld.global.b32 %0, [%1];" : "=&r"(rB0) : "l"(b0) : "memory");
+                 asm volatile("ld.global.b32 %0, [%1];" : "=&r"(rB1) : "l"(b1) : "memory");
+                 asm volatile("ld.global.b32 %0, [%1];" : "=&r"(rB2) : "l"(b2) : "memory");
+                 asm volatile("ld.global.b32 %0, [%1];" : "=&r"(rB3) : "l"(b3) : "memory");
+  
+                 asm volatile("add.f16x2 %0, %1, %2;" : "=&r"(rC0) : "r"(rA0), "r"(rB0));
+                 asm volatile("add.f16x2 %0, %1, %2;" : "=&r"(rC1) : "r"(rA1), "r"(rB1));
+                 asm volatile("add.f16x2 %0, %1, %2;" : "=&r"(rC2) : "r"(rA2), "r"(rB2));
+                 asm volatile("add.f16x2 %0, %1, %2;" : "=&r"(rC3) : "r"(rA3), "r"(rB3));
+  
+                 asm volatile("st.global.b32 [%0], %1;" :: "l"(c0), "r"(rC0) : "memory");
+                 asm volatile("st.global.b32 [%0], %1;" :: "l"(c1), "r"(rC1) : "memory");
+                 asm volatile("st.global.b32 [%0], %1;" :: "l"(c2), "r"(rC2) : "memory");
+                 asm volatile("st.global.b32 [%0], %1;" :: "l"(c3), "r"(rC3) : "memory");
+             }
+         } else {
+             // Tail: scalar half add for remaining < VEC elements
+             for (int64_t t = 0; t < remaining; ++t) {
+                 C[idx + t] = __hadd(A[idx + t], B[idx + t]);
+             }
          }
      }
  }
   
- torch::Tensor vectoradd_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor out) {
-     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* out_ptr = reinterpret_cast<half*>(out.data_ptr<at::Half>());
+ void vec_add_ptx(torch::Tensor A, torch::Tensor B, torch::Tensor C) {
+     TORCH_CHECK(A.is_cuda() && B.is_cuda() && C.is_cuda(), "Tensors must be CUDA");
+     TORCH_CHECK(A.scalar_type() == at::kHalf && B.scalar_type() == at::kHalf && C.scalar_type() == at::kHalf, "Tensors must be float16 (half)");
+     TORCH_CHECK(A.is_contiguous() && B.is_contiguous() && C.is_contiguous(), "Tensors must be contiguous");
+     TORCH_CHECK(A.numel() == B.numel() && A.numel() == C.numel(), "Size mismatch");
   
-     const int threads = 256;
-     const int blocks = (n + threads * 16 - 1) / (threads * 16);
+     c10::cuda::CUDAGuard guard(A.device());
+     const int64_t N = A.numel();
+     if (N == 0) return;
   
-     vectoradd_kernel<<<blocks, threads>>>(a_ptr, b_ptr, out_ptr, n);
-     return out;
+     const int threads = 256;      // good balance for SM80
+     const int VEC = 8;            // 8 fp16 elems (16B) per thread for 128-bit transactions
+     const int64_t elems_per_block = int64_t(threads) * VEC;
+     int blocks = static_cast<int>((N + elems_per_block - 1) / elems_per_block);
+     blocks = std::min(blocks, 65535); // cap to keep scheduler happy
+  
+     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+     vec_add_half2_ptx<VEC><<<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*>(C.data_ptr<at::Half>()),
+         N
+     );
+     AT_CUDA_CHECK(cudaGetLastError());
  }
- """
   
- cpp_source = "torch::Tensor vectoradd_cuda(torch::Tensor, torch::Tensor, torch::Tensor);"
+ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+     m.def("vec_add_ptx", &vec_add_ptx, "Vector add using inline PTX (half)");
+ }
+ '''
   
- _module = None
+ _ext_mod = None
   
- def _get_module():
-     global _module
-     if _module is None:
-         _module = load_inline(
-             name='vectoradd_coalesced',
-             cpp_sources=cpp_source,
-             cuda_sources=cuda_source,
-             functions=['vectoradd_cuda'],
-             extra_cuda_cflags=['-O3', '-use_fast_math', '-arch=sm_80', '-maxrregcount=64'],
-             verbose=False
+ def _get_ext_mod():
+     global _ext_mod
+     if _ext_mod is None:
+         _ext_mod = load_inline(
+             name="vecadd_half_ptx_sm80",
+             cpp_sources="",
+             cuda_sources=_cuda_src,
+             extra_cuda_cflags=[
+                 "-O3",
+                 "--use_fast_math",
+                 "-U__CUDA_NO_HALF_OPERATORS__",
+                 "-U__CUDA_NO_HALF_CONVERSIONS__",
+                 "-gencode=arch=compute_80,code=sm_80",
+             ],
+             verbose=False,
          )
-     return _module
+     return _ext_mod
   
  def custom_kernel(data: input_t) -> output_t:
      A, B, output = data
-     _get_module().vectoradd_cuda(A, B, output)
+  
+     # Fast path: CUDA half, contiguous, same numel
+     if (
+         A.is_cuda and B.is_cuda and output.is_cuda and
+         A.dtype == torch.float16 and B.dtype == torch.float16 and output.dtype == torch.float16 and
+         A.is_contiguous() and B.is_contiguous() and output.is_contiguous() and
+         A.numel() == B.numel() == output.numel()
+     ):
+         mod = _get_ext_mod()
+         mod.vec_add_ptx(A, B, output)
+         return output
+  
+     # Safe fallback without allocating a new output tensor
+     output.copy_(A)
+     output.add_(B)
      return output