submission 66199

  ⋯ 1 unchanged lines
  import triton
  import triton.language as tl
  
+ 
+ @triton.autotune(
+     configs=[
+         triton.Config({'BLOCK_SIZE': 256}, num_warps=2, num_stages=2),
+         triton.Config({'BLOCK_SIZE': 512}, num_warps=2, num_stages=2),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=4, num_stages=2),
+         triton.Config({'BLOCK_SIZE': 2048}, num_warps=4, num_stages=2),
+         triton.Config({'BLOCK_SIZE': 2048}, num_warps=8, num_stages=2),
+         triton.Config({'BLOCK_SIZE': 4096}, num_warps=8, num_stages=2),
+         triton.Config({'BLOCK_SIZE': 4096}, num_warps=8, num_stages=3),
+     ],
+     key=['N'],
+ )
  @triton.jit
  def vectoradd_kernel(
-     a_ptr,  # Pointer to first input tensor
-     b_ptr,  # Pointer to second input tensor
-     c_ptr,  # Pointer to output tensor
-     N,      # Total number of elements
-     BLOCK_SIZE: tl.constexpr,  # Number of elements per block
+     a_ptr,
+     b_ptr,
+     c_ptr,
+     N,
+     BLOCK_SIZE: tl.constexpr,
  ):
      """
-     Optimized float16 vector addition kernel.
-     Each program processes BLOCK_SIZE elements.
+     Ultra-optimized float16 vector addition for B200.
+     Uses autotuning for optimal configuration.
      """
-     # Get program ID
      pid = tl.program_id(0)
-     
-     # Compute block start offset
      block_start = pid * BLOCK_SIZE
-     
-     # Generate offsets for this block
      offsets = block_start + tl.arange(0, BLOCK_SIZE)
-     
-     # Create mask for boundary checking (handles non-multiple of BLOCK_SIZE)
      mask = offsets < N
      
-     # Load data with vectorized operations (float16)
-     # Use eviction policy for better cache utilization
-     a = tl.load(a_ptr + offsets, mask=mask, other=0.0, eviction_policy='evict_first')
-     b = tl.load(b_ptr + offsets, mask=mask, other=0.0, eviction_policy='evict_first')
+     # Load with cache eviction hints optimized for B200
+     a = tl.load(a_ptr + offsets, mask=mask, other=0.0, eviction_policy='evict_last')
+     b = tl.load(b_ptr + offsets, mask=mask, other=0.0, eviction_policy='evict_last')
      
      # Perform addition
      c = a + b
      
-     # Store result with vectorized write
-     tl.store(c_ptr + offsets, c, mask=mask)
+     # Store with write-back policy
+     tl.store(c_ptr + offsets, c, mask=mask, eviction_policy='evict_last')
  
  
  def custom_kernel(data):
      """
-     Main kernel function that launches the Triton kernel.
-     
-     Args:
-         data: Tuple of (A, B, C) where A and B are input tensors,
-               and C is the pre-allocated output tensor.
-     
-     Returns:
-         C: Output tensor containing element-wise sum
+     Main kernel with autotuning for peak B200 performance.
      """
      A, B, C = data
      
-     # Ensure inputs are contiguous for optimal memory access
+     # Ensure contiguous layout
      A = A.contiguous()
      B = B.contiguous()
+     C = C.contiguous()
      
-     # Get tensor dimensions
      N = A.numel()
      
-     # Choose optimal block size based on GPU architecture
-     # Larger block sizes for better memory coalescing and fewer kernel launches
-     # Powers of 2 work best for GPU memory systems
-     BLOCK_SIZE = 1024  # Optimal for most modern GPUs (A100, H100, B200)
-     
-     # Calculate grid size (number of programs to launch)
+     # Grid configuration - will be optimized by autotune
      grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']),)
      
-     # Launch kernel
-     vectoradd_kernel[grid](
-         A, B, C, N,
-         BLOCK_SIZE=BLOCK_SIZE,
-     )
+     # Launch autotuned kernel
+     vectoradd_kernel[grid](A, B, C, N)
      
      return C