#!POPCORN leaderboard amd-identity
from task import input_t, output_t
import torch
import triton.language as tl
import triton


@triton.jit
def copy_kern(dst_ptr, src_ptr, n_elems, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    tl.store(dst_ptr + offsets, tl.load(src_ptr + offsets))


def copy(dst, src):
    numel = dst.numel()

    def grid(meta):
        return (triton.cdiv(numel, meta["BLOCK_SIZE"]),)

    copy_kern[grid](dst, src, numel, BLOCK_SIZE=1024)


def custom_kernel(data: input_t) -> output_t:
    input, output = data
    copy(output, input)
    # output[...] = input
    return output


if __name__ == "__main__":
    x = torch.randn(1000, 1000, device="cuda")
    y = torch.randn(1000, 1000, device="cuda")
    custom_kernel((x, y))
    assert (x == y).all()

python eval.py benchmark /tmp/tmpdzjfwsmc