#!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 torch
import triton
import triton.language as tl

@triton.jit
def vector_copy_kernel(input_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)  # Each program ID handles one block of data
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)  # Compute thread indices
    mask = offsets < n_elements  # Don't go out of bounds

    # Load input, store to output
    input_vals = tl.load(input_ptr + offsets, mask=mask)
    tl.store(output_ptr + offsets, input_vals, 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 = input.numel()

    # Ensure input/output are contiguous and on the right device
    assert input.is_cuda and output.is_cuda and input.dtype == torch.float16

    # Launch Triton kernel
    BLOCK_SIZE = 1024
    grid = lambda META: (triton.cdiv(n, META['BLOCK_SIZE']),)

    vector_copy_kernel[grid](input, output, n, BLOCK_SIZE=BLOCK_SIZE)
    return output


python eval.py benchmark /tmp/tmpp5s6094_