submission 639715

  # kernel.py
- """
- Float16 vector/matrix addition implemented in Triton.
+ # Fused pipeline: (1) load A + load B -> (2) elementwise add -> (3) store C
+ # Everything is done in a single Triton kernel (no unfused stages needed).
+ #
+ # Fix for timeout: remove @triton.autotune keyed on n_elements, which can trigger
+ # multiple compilations/benchmarks for each different input size and exceed the
+ # test timeout. Use a single compiled configuration instead.
   
- Fused pipeline (single pass):
-   1) tl.load A
-   2) tl.load B
-   3) elementwise add in-kernel
-   4) tl.store C
-  
- Wrapper performs only validation/allocation/grid setup/launch (no PyTorch math).
- """
-  
- from __future__ import annotations
-  
  import torch
  import triton
  import triton.language as tl
   
   
- @triton.autotune(
-     configs=[
-         triton.Config({"BLOCK_SIZE": 256}, num_warps=4),
-         triton.Config({"BLOCK_SIZE": 512}, num_warps=4),
-         triton.Config({"BLOCK_SIZE": 1024}, num_warps=8),
-     ],
-     key=["n_elements"],
- )
  @triton.jit
- def _vec_add_kernel(
-     a_ptr, b_ptr, c_ptr,
-     n_elements: tl.int32,
+ def _add_f16_kernel(
+     a_ptr,
+     b_ptr,
+     c_ptr,
+     n_elements,
      BLOCK_SIZE: tl.constexpr,
  ):
      pid = tl.program_id(axis=0)
-     block_start = pid * BLOCK_SIZE
-     offsets = block_start + tl.arange(0, BLOCK_SIZE)
-     mask = offsets < n_elements
+     offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+     mask = offs < n_elements
   
-     a = tl.load(a_ptr + offsets, mask=mask, other=0.0)
-     b = tl.load(b_ptr + offsets, mask=mask, other=0.0)
-     c = a + b
-     tl.store(c_ptr + offsets, c, mask=mask)
+     # Contiguous 1D flatten => naturally coalesced.
+     a = tl.load(a_ptr + offs, mask=mask, other=0.0)
+     b = tl.load(b_ptr + offs, mask=mask, other=0.0)
+     tl.store(c_ptr + offs, a + b, mask=mask)
   
   
- def kernel_function(*args):
+ def kernel_function(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor):
      """
-     Supported call patterns (as used by the test harness):
-       - kernel_function(A, B) -> returns new C
-       - kernel_function(A, B, C) -> writes into C (and may return C/None)
-       - kernel_function((A, B)) / kernel_function((A, B, C))
+     Elementwise float16 add for two (N, N) CUDA tensors: C = A + B.
   
-     Inputs:
-       A, B: (N, N) float16 CUDA contiguous
-       C:    optional output buffer, same shape/dtype/device
+     Wrapper responsibilities only: validate/allocate/launch. No PyTorch math ops.
      """
-     # Unpack tuple-packed variants
-     if len(args) == 1 and isinstance(args[0], (tuple, list)):
-         args = tuple(args[0])
+     assert isinstance(A, torch.Tensor) and isinstance(B, torch.Tensor) and isinstance(C, torch.Tensor)
+     assert A.is_cuda and B.is_cuda and C.is_cuda
+     assert A.dtype == torch.float16 and B.dtype == torch.float16 and C.dtype == torch.float16
+     assert A.shape == B.shape == C.shape
+     assert A.is_contiguous() and B.is_contiguous() and C.is_contiguous()
   
-     if len(args) not in (2, 3):
-         raise TypeError(f"kernel_function expected 2 or 3 arguments, got {len(args)}")
-  
-     A = args[0]
-     B = args[1]
-     C = args[2] if len(args) == 3 else None
-  
-     if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor):
-         raise TypeError("A and B must be torch.Tensors")
-     if A.device.type != "cuda" or B.device.type != "cuda":
-         raise ValueError("A and B must be CUDA tensors")
-     if A.dtype != torch.float16 or B.dtype != torch.float16:
-         raise ValueError("A and B must be float16")
-     if A.shape != B.shape:
-         raise ValueError(f"Shape mismatch: A.shape={tuple(A.shape)} B.shape={tuple(B.shape)}")
-     if not A.is_contiguous() or not B.is_contiguous():
-         raise ValueError("A and B must be contiguous")
-     if A.ndim != 2 or A.shape[0] != A.shape[1]:
-         raise ValueError(f"Expected A and B to have shape (N, N), got {tuple(A.shape)}")
-  
-     if C is None:
-         C = torch.empty_like(A)
-     else:
-         if not isinstance(C, torch.Tensor):
-             raise TypeError("C must be a torch.Tensor when provided")
-         if C.device != A.device:
-             raise ValueError("C must be on the same device as A")
-         if C.dtype != torch.float16:
-             raise ValueError("C must be float16")
-         if C.shape != A.shape:
-             raise ValueError("C must have the same shape as A")
-         if not C.is_contiguous():
-             raise ValueError("C must be contiguous")
-  
      n_elements = A.numel()
   
-     def grid(meta):
-         return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+     # Single stable configuration (fast compile, avoids autotune timeouts).
+     BLOCK_SIZE = 1024
+     grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
   
-     _vec_add_kernel[grid](
+     _add_f16_kernel[grid](
          A, B, C,
-         n_elements=n_elements,
+         n_elements,
+         BLOCK_SIZE=BLOCK_SIZE,
+         num_warps=8,
      )
      return C