#!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)


@triton.jit
def fast_vector_copy_kernel(input_ptr, output_ptr, num_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)

    # Each thread handles 4 float16s at a time => 64-bit vector copy
    base_offsets = pid * BLOCK_SIZE * 4 + tl.arange(0, BLOCK_SIZE) * 4  # [BLOCK_SIZE]
    inner_offsets = tl.arange(0, 4)  # [4]
    full_offsets = base_offsets[:, None] + inner_offsets[None, :]       # [BLOCK_SIZE, 4]

    mask = full_offsets < num_elements
    in_vals = tl.load(input_ptr + full_offsets, mask=mask)
    tl.store(output_ptr + full_offsets, in_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

    numel = input.numel()

    BLOCK_SIZE = 1024  # number of threads (each loads/stores 4 elements)
    grid = lambda META: (triton.cdiv(numel, BLOCK_SIZE * 4),)

    fast_vector_copy_kernel[grid](
        input, output, numel,
        BLOCK_SIZE=BLOCK_SIZE
    )

    return output


python eval.py test /tmp/tmp835a99fx