submission 612308

  import os
  os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
  import torch
- from torch.utils.cpp_extension import load_inline
+ import triton
+ import triton.language as tl
  from task import input_t, output_t
   
- cuda_source = r"""
- #include <torch/extension.h>
- #include <cub/cub.cuh>
- #include <cuda_runtime.h>
+ BLOCK_SIZE = 1024
   
- static void* d_temp_storage = nullptr;
- static size_t temp_storage_bytes = 0;
- static int last_n = 0;
+ @triton.jit
+ def _reduce_kernel(input_ptr, block_sums_ptr, n, BLOCK_SIZE: tl.constexpr):
+ 	pid = tl.program_id(0)
+ 	offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+ 	mask = offsets < n
+ 	x = tl.load(input_ptr + offsets, mask=mask, other=0.0)
+ 	block_sum = tl.sum(x)
+ 	tl.store(block_sums_ptr + pid, block_sum)
   
- void inclusive_scan(torch::Tensor input, torch::Tensor output) {
-     int n = input.numel();
-     float* d_in = input.data_ptr<float>();
-     float* d_out = output.data_ptr<float>();
+ @triton.jit
+ def _scan_and_add_kernel(input_ptr, output_ptr, prefix_sums_ptr, n, BLOCK_SIZE: tl.constexpr):
+ 	pid = tl.program_id(0)
+ 	offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+ 	mask = offsets < n
+ 	x = tl.load(input_ptr + offsets, mask=mask, other=0.0)
+ 	scanned = tl.cumsum(x)
+ 	if pid > 0:
+ 		prefix = tl.load(prefix_sums_ptr + pid - 1)
+ 	else:
+ 		prefix = 0.0
+ 	scanned = scanned + prefix
+ 	tl.store(output_ptr + offsets, scanned, mask=mask)
   
-     if (n != last_n) {
-         if (d_temp_storage) cudaFree(d_temp_storage);
-         temp_storage_bytes = 0;
-         cub::DeviceScan::InclusiveSum(nullptr, temp_storage_bytes, d_in, d_out, n);
-         cudaMalloc(&d_temp_storage, temp_storage_bytes);
-         last_n = n;
-     }
+ # Pre-allocate buffers for the target size
+ _block_sums = None
+ _prefix_sums = None
   
-     cub::DeviceScan::InclusiveSum(d_temp_storage, temp_storage_bytes, d_in, d_out, n);
- }
- """
+ def custom_kernel(data: input_t) -> output_t:
+ 	global _block_sums, _prefix_sums
+ 	inp, out = data
+ 	n = inp.numel()
+ 	nb = (n + BLOCK_SIZE - 1) // BLOCK_SIZE
   
- cpp_source = r"""
- void inclusive_scan(torch::Tensor input, torch::Tensor output);
- """
+ 	if _block_sums is None or _block_sums.numel() < nb:
+ 		_block_sums = torch.empty(nb, device="cuda", dtype=torch.float32)
+ 		_prefix_sums = torch.empty(nb, device="cuda", dtype=torch.float32)
   
- module = load_inline(
-     name="cub_scan",
-     cpp_sources=[cpp_source],
-     cuda_sources=[cuda_source],
-     functions=["inclusive_scan"],
-     extra_cuda_cflags=["-O3", "--use_fast_math"],
-     verbose=False,
- )
-  
-  
- def custom_kernel(data: input_t) -> output_t:
-     data, output = data
-     module.inclusive_scan(data, output)
-     return output
+ 	_reduce_kernel[(nb,)](inp, _block_sums, n, BLOCK_SIZE=BLOCK_SIZE, num_warps=4)
+ 	torch.cumsum(_block_sums[:nb], dim=0, out=_prefix_sums[:nb])
+ 	_scan_and_add_kernel[(nb,)](inp, out, _prefix_sums, n, BLOCK_SIZE=BLOCK_SIZE, num_warps=4)
+ 	return out