#!POPCORN leaderboard amd-identity

# This is a submission template for popcorn leaderboard 'amd-identity'.
# Your task is as follows:
# > This task is purely for testing the submission system. There will be *no* points.
# > 
# > > Input: (input_tensor, output_tensor)
# > >  - input_tensor: Input data
# > >  - output_tensor: Pre-allocated empty tensor of the same shape as `input_tensor`
# > > Output: Should return `output_tensor` after it has been filled with the values from `input_tensor`.`
# The deadline for this leaderboard is 2025-06-02 00:00:00+00:00

# You can automatically route this file to specific GPUs by adding a line
# `#!POPCORN gpus <GPUs>` to the header of this file.
# Happy hacking!

from task import input_t, output_t
import triton
import triton.language as tl

@triton.jit
def copy_kernel(
    x_ptr,
    y_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
    
    x = tl.load(x_ptr + offsets, mask=mask)
    tl.store(y_ptr + offsets, x, mask=mask)

def custom_kernel(data: input_t) -> output_t:
    """
    Copies the contents of `input` into `output`
    Args:
        data: tuple of (input, output) tensors

    Returns: output tensor
    """
    input, output = data
    n_elements = input.numel()
    
    # Calculate grid size
    BLOCK_SIZE = 1024
    grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
    
    # Launch kernel
    copy_kernel[grid](
        input,
        output,
        n_elements,
        BLOCK_SIZE=BLOCK_SIZE,
    )
    
    return output


python eval.py benchmark /tmp/tmp_qpzrnar