submission 102316

  import os
  # Unset CUDA_VISIBLE_DEVICES before importing torch/triton
- # This works around SLURM setting it to a value that prevents PyTorch from seeing GPUs
  if 'CUDA_VISIBLE_DEVICES' in os.environ:
      del os.environ['CUDA_VISIBLE_DEVICES']
   
- import triton.testing
- import math
- import logging
-  
+ import torch
  import triton
  import triton.language as tl
- import torch
   
- torch.cuda.empty_cache()
- print(f"CUDA available: {torch.cuda.is_available()}")
- print(f"Device count: {torch.cuda.device_count()}")
- import torch.nn.functional as F
   
  @triton.jit
- def naive_prefixsum_kernel(
+ def prefixsum_kernel(
      input_ptr,
      output_ptr,
      n_elements,
      BLOCK_SIZE: tl.constexpr,
  ):
+     """Compute prefix sum within each block."""
      pid = tl.program_id(axis=0)
      block_start = pid * BLOCK_SIZE
      offsets = block_start + tl.arange(0, BLOCK_SIZE)
      mask = offsets < n_elements
   
-     # Load input
+     # Load data from input
      x = tl.load(input_ptr + offsets, mask=mask, other=0.0)
   
-     # Single-thread accumulator
-     # Only lane 0 computes the prefix sum sequentially.
-     acc = tl.zeros((BLOCK_SIZE,), dtype=x.dtype)
+     # Compute prefix sum within block
+     x = tl.cumsum(x, axis=0)
   
-     # get the lane id within the program
-     lane_id = tl.arange(0, BLOCK_SIZE)
+     # Store the result
+     tl.store(output_ptr + offsets, x, mask=mask)
   
-     # Run prefix sum sequentially on lane 0
-     if tl.program_id(axis=1) == 0:  # always true (1D grid) but needed for Triton scoping
-         # Only lane 0 executes the loop
-         if lane_id == 0:
-             running = 0.0
-             for i in range(BLOCK_SIZE):
-                 running += x[i]
-                 acc[i] = running
   
-     # Broadcast results: each lane loads its own index from acc
-     out = acc[lane_id]
-  
-     tl.store(output_ptr + offsets, out, mask=mask)
-  
-  
-  
-  
- # --------- kernel ----------------
  @triton.jit
- def prefixsum_kernel(
-     input_ptr,      # pointer to input array
-     output_ptr,     # pointer to output array
-     n_elements,     # total number of elements
-     BLOCK_SIZE: tl.constexpr,  # block size for processing
+ def extract_block_sums_kernel(
+     data_ptr,
+     block_sums_ptr,
+     n_elements,
+     BLOCK_SIZE: tl.constexpr,
  ):
+     """Extract the last element of each block (the block sum)."""
      pid = tl.program_id(axis=0)
-     block_start = pid * BLOCK_SIZE
-     offsets = block_start + tl.arange(0, BLOCK_SIZE)
-     mask = offsets < n_elements
   
-     # Load data from input, with padding for elements out of bounds
-     x = tl.load(input_ptr + offsets, mask=mask, other=0.0)
+     # Last index in this block
+     last_idx = min((pid + 1) * BLOCK_SIZE - 1, n_elements - 1)
   
-     # Perform parallel prefix sum within the block using tl.cumsum.
-     # cumsum is the built-in prefix sum operation in Triton.
-     x = tl.cumsum(x, axis=0)
+     # Load the last element (sum of entire block)
+     block_sum = tl.load(data_ptr + last_idx)
   
-     # Store the result to output
-     tl.store(output_ptr + offsets, x, mask=mask)
+     # Store to block_sums array
+     tl.store(block_sums_ptr + pid, block_sum)
   
+  
  @triton.jit
- def add_block_sums_kernel(
+ def add_block_offsets_kernel(
      output_ptr,
-     block_sums_ptr,
+     block_offsets_ptr,
      n_elements,
      BLOCK_SIZE: tl.constexpr,
  ):
-     """Add the cumulative sum from previous blocks to each element."""
+     """Add the offset from previous blocks to each element."""
      pid = tl.program_id(axis=0)
-     
+  
      if pid == 0:
-         return  # First block doesn't need adjustment
-     
+         return  # First block has no offset
+  
      block_start = pid * BLOCK_SIZE
      offsets = block_start + tl.arange(0, BLOCK_SIZE)
      mask = offsets < n_elements
-     
-     # Load the cumulative sum from all previous blocks
-     prev_sum = tl.load(block_sums_ptr + pid - 1)
-     
-     # Load current values, add previous sum, and store back
+  
+     # Load the cumulative offset for this block
+     offset = tl.load(block_offsets_ptr + pid - 1)
+  
+     # Load current values, add offset, and store back
      x = tl.load(output_ptr + offsets, mask=mask)
-     x = x + prev_sum
+     x = x + offset
      tl.store(output_ptr + offsets, x, mask=mask)
   
   
- # --------- wrapper function ----------------
- def prefixsum(data):
+ def prefixsum_recursive(input_tensor, output_tensor, block_size=1024):
      """
-     Wrapper function to launch the prefix sum kernel.
+     Recursive GPU-based prefix sum.
+  
      Args:
-         data: Tuple of (input_tensor, output_tensor)
+         input_tensor: Input tensor (1D)
+         output_tensor: Output tensor (1D, same size as input)
+         block_size: Size of each block for parallel processing
+  
      Returns:
-         output_tensor: Tensor containing the inclusive prefix sum
+         output_tensor: Result of prefix sum
      """
-     input_tensor, output_tensor = data
      n_elements = input_tensor.numel()
-     
-     # Choose block size
-     BLOCK_SIZE = 4096
-     
-     # Calculate grid size
-     num_blocks = triton.cdiv(n_elements, BLOCK_SIZE)
+     num_blocks = triton.cdiv(n_elements, block_size)
+  
+     # Base case: single block
+     if num_blocks == 1:
+         grid = (1,)
+         prefixsum_kernel[grid](
+             input_tensor,
+             output_tensor,
+             n_elements,
+             BLOCK_SIZE=block_size,
+         )
+         return output_tensor
+  
+     # Recursive case: multiple blocks
      grid = (num_blocks,)
-     
-     # Pass 1: Compute prefix sum within each block
+  
+     # Step 1: Compute prefix sum within each block
      prefixsum_kernel[grid](
          input_tensor,
          output_tensor,
          n_elements,
-         BLOCK_SIZE=BLOCK_SIZE,
+         BLOCK_SIZE=block_size,
      )
   
-     # If we have multiple blocks, we need to propagate sums across blocks
-     if num_blocks > 1:
-         # Extract the last element of each block (the block sums)
-         block_sums = torch.zeros(num_blocks, dtype=input_tensor.dtype, device=input_tensor.device)
-         for i in range(num_blocks):
-             last_idx = min((i + 1) * BLOCK_SIZE - 1, n_elements - 1)
-             block_sums[i] = output_tensor[last_idx]
+     # Step 2: Extract block sums (last element of each block)
+     block_sums = torch.empty(num_blocks, dtype=input_tensor.dtype, device=input_tensor.device)
+     extract_block_sums_kernel[grid](
+         output_tensor,
+         block_sums,
+         n_elements,
+         BLOCK_SIZE=block_size,
+     )
   
-         # Compute prefix sum of block sums
-         block_prefix = torch.cumsum(block_sums, dim=0)
+     # Step 3: Recursively compute prefix sum of block sums
+     block_offsets = torch.empty_like(block_sums)
+     prefixsum_recursive(block_sums, block_offsets, block_size=block_size)
   
-         # Pass 2: Add accumulated sums from previous blocks
-         add_block_sums_kernel[grid](
-             output_tensor,
-             block_prefix,
-             n_elements,
-             BLOCK_SIZE=BLOCK_SIZE,
-         )
+     # Step 4: Add block offsets to all elements
+     add_block_offsets_kernel[grid](
+         output_tensor,
+         block_offsets,
+         n_elements,
+         BLOCK_SIZE=block_size,
+     )
   
      return output_tensor
   
   
- # Alias for submission system
- custom_kernel = prefixsum
   No newline at end of file
+ def custom_kernel(data, block_size=1024):
+     """
+     Optimized prefix sum using fully GPU-based recursive algorithm.
+  
+     Args:
+         data: Tuple of (input_tensor, output_tensor)
+         block_size: Block size for processing (default 1024)
+                    Larger blocks = fewer recursion levels but less parallelism per block
+                    Smaller blocks = more recursion but better intra-block parallelism
+  
+     Returns:
+         output_tensor: Result of prefix sum
+     """
+     input_tensor, output_tensor = data
+  
+     # Auto-tune block size based on input size if desired
+     n = input_tensor.numel()
+     if block_size == 'auto':
+         # Heuristic: use larger blocks for larger inputs
+         if n < 1024:
+             block_size = 256
+         elif n < 16384:
+             block_size = 1024
+         elif n < 262144:
+             block_size = 2048
+         else:
+             block_size = 4096
+  
+     return prefixsum_recursive(input_tensor, output_tensor, block_size=block_size)
+  
+  
+ # Backward compatibility
+ prefixsum = custom_kernel