submission 560058

  # kernel.py
  """
- Triton vector-add kernel (float16) for 2D tensors.
+ Float16 vector/matrix addition implemented in Triton.
   
- Fused stages (single pass):
-   1) tl.load(A) + tl.load(B)
-   2) elementwise add in-kernel
-   3) tl.store(C)
+ Fused pipeline (single pass):
+   1) tl.load A
+   2) tl.load B
+   3) elementwise add in-kernel
+   4) tl.store C
   
- No unfused fallback is needed because the whole pipeline is a single elementwise op.
+ 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,
-                     BLOCK_SIZE: tl.constexpr):
+ def _vec_add_kernel(
+     a_ptr, b_ptr, c_ptr,
+     n_elements: tl.int32,
+     BLOCK_SIZE: tl.constexpr,
+ ):
      pid = tl.program_id(axis=0)
-     offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-     mask = offs < n_elements
+     block_start = pid * BLOCK_SIZE
+     offsets = block_start + tl.arange(0, BLOCK_SIZE)
+     mask = offsets < n_elements
   
-     a = tl.load(a_ptr + offs, mask=mask, other=0.0)
-     b = tl.load(b_ptr + offs, mask=mask, other=0.0)
+     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 + offs, c, mask=mask)
+     tl.store(c_ptr + offsets, c, mask=mask)
   
   
  def kernel_function(*args):
      """
-     Wrapper that validates inputs, allocates output if needed, and launches Triton.
+     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))
   
-     Supported call patterns (as used by the tests):
-       - kernel_function(A, B) -> out
-       - kernel_function(A, B, C) -> out (writes into C)
-       - kernel_function((A, B)) -> out
-       - kernel_function((A, B, C)) -> out (writes into C)
+     Inputs:
+       A, B: (N, N) float16 CUDA contiguous
+       C:    optional output buffer, same shape/dtype/device
      """
-     # Unpack tuple-style inputs
+     # Unpack tuple-packed variants
      if len(args) == 1 and isinstance(args[0], (tuple, list)):
          args = tuple(args[0])
   
      if len(args) not in (2, 3):
          raise TypeError(f"kernel_function expected 2 or 3 arguments, got {len(args)}")
   
-     A, B = args[0], args[1]
+     A = args[0]
+     B = args[1]
      C = args[2] if len(args) == 3 else None
   
-     if not (isinstance(A, torch.Tensor) and isinstance(B, torch.Tensor)):
-         raise TypeError("A and B must be torch.Tensor")
+     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 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)} vs B.shape={tuple(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)
  ⋯ 3 unchanged lines
          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 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()
-     BLOCK_SIZE = 1024
-     grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
   
+     def grid(meta):
+         return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+  
      _vec_add_kernel[grid](
          A, B, C,
-         n_elements,
-         BLOCK_SIZE=BLOCK_SIZE,
-         num_warps=8,
+         n_elements=n_elements,
      )
      return C