# Copyright (c) Tile-AI Organization.
# Licensed under the MIT License.

# download tilelang-rocm
from pip._internal import main as pip_main
pip_main(["install","tilelang-rocm"])
pip_main(["install","apache-tvm"])
from task import input_t, output_t
from reference import generate_input, ref_kernel
import torch
import tilelang
import tilelang.language as T
from tilelang.intrinsics import make_mfma_swizzle_layout as make_swizzle_layout
from tilelang.intrinsics.mfma_macro_generator import (
    MatrixCoreIntrinEmitter,)
from tilelang.intrinsics.mfma_layout import thread_id_shared_access_64x4_to_16x16_layout_C_n_m
import itertools
import tilelang as tl
import tilelang.language as T
from tilelang.autotuner import AutoTuner
# from tilelang.carver.template import MatmulTemplate
# from tilelang.carver.arch import CDNA
# from tilelang.carver.roller.rasterization import NoRasterization


def ref_program(A, B, A_scale, B_scale):
    C = torch.zeros((A.shape[0], B.shape[0]), dtype=torch.bfloat16, device="cuda")
    return ref_kernel((A, B, A_scale, B_scale, C))

def get_configs(M, N, K, with_roller=False, topk=20):
    block_col_warps = [2]
    block_row_warps = [2]
    num_stages = [0, 1, 2]
    thread_num = [128]
    enable_rasterization = [True, False]
    _configs = list(
        itertools.product(
            block_col_warps,
            block_row_warps,
            num_stages,
            thread_num,
            enable_rasterization,
        ))

    configs = [
        {
            "block_col_warps": c[0],
            "block_row_warps": c[1],
            "num_stages": c[2],
            "thread_num": c[3],
            "enable_rasterization": c[4],
        } for c in _configs
    ]
    return configs

def get_best_config(M, N, K, with_roller=False):

    def kernel(
        block_col_warps,
        block_row_warps,
        num_stages,
        thread_num,
        enable_rasterization,
        in_dtype="float8_e4m3fnuz",
        out_dtype="bfloat16",
        accum_dtype="float32",
        a_transposed=False,
        b_transposed=True,
    ):

        micro_size_x = micro_size_y = micro_size_k = 16

        if in_dtype in {"float8_e4m3fnuz", "int8"}:
            micro_size_k = 32
        
        k_pack = 2
        warp_row_tiles = 64
        warp_col_tiles = 64
        chunk = 128
        shared_scope = "shared"
        cache_write_shared = False

        block_M = 128
        block_N = 128
        block_K = chunk

        scale_n =  (N + block_N - 1) // block_N
        scale_k =  (K + block_K - 1) // block_K

        A_shape = (K, M) if a_transposed else (M, K)
        B_shape = (N, K) if b_transposed else (K, N)
        A_scale_shape = (M, scale_k)
        B_scale_shape = (scale_n, scale_k)

        A_shared_shape = (block_K, block_M) if a_transposed else (block_M, block_K)
        B_shared_shape = (block_N, block_K) if b_transposed else (block_K, block_N)
        C_shared_shape = (
            block_N // micro_size_y,
            micro_size_x,
            micro_size_y,
        )

        warp_size = 64
        threads = warp_size * (block_row_warps * block_col_warps)
        local_size_a = (k_pack * micro_size_x * micro_size_k) // warp_size
        local_size_b = (k_pack * micro_size_y * micro_size_k) // warp_size
        local_size_c = (micro_size_x * micro_size_y) // warp_size
        warp_rows = warp_row_tiles // micro_size_x
        warp_cols = warp_col_tiles // micro_size_y

        # MMA Wrapper to Auto Generate Code for MMA
        mfma_emitter = MatrixCoreIntrinEmitter(
            a_dtype=in_dtype,
            b_dtype=in_dtype,
            accum_dtype=accum_dtype,
            a_transposed=a_transposed,
            b_transposed=b_transposed,
            block_row_warps=block_row_warps,
            block_col_warps=block_col_warps,
            warp_row_tiles=warp_row_tiles,
            warp_col_tiles=warp_col_tiles,
            chunk=chunk,
            k_pack=k_pack,
        )

        @T.prim_func
        def main(
                A: T.Tensor(A_shape, in_dtype),
                B: T.Tensor(B_shape, in_dtype),
                A_scale: T.Tensor(A_scale_shape, accum_dtype),
                B_scale: T.Tensor(B_scale_shape, accum_dtype),
                C: T.Tensor((M, N), out_dtype),
        ):
            with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):

                A_shared = T.alloc_shared(A_shared_shape, in_dtype, scope=shared_scope)
                B_shared = T.alloc_shared(B_shared_shape, in_dtype, scope=shared_scope)
                C_shared = T.alloc_shared(C_shared_shape, out_dtype, scope=shared_scope)
                A_local = T.alloc_local((warp_rows * local_size_a), in_dtype)
                B_local = T.alloc_local((warp_cols * local_size_b), in_dtype)
                C_local = T.alloc_local((warp_rows * warp_cols * local_size_c), accum_dtype)
                C_local_accum = T.alloc_local((warp_rows * warp_cols * local_size_c), accum_dtype)

                T.annotate_layout({
                    A_shared: make_swizzle_layout(A_shared),
                    B_shared: make_swizzle_layout(B_shared),
                })

                # Improve L2 Cache
                T.use_swizzle(panel_size=10, enable=enable_rasterization)
                
                thread_binding = T.get_thread_binding()
                tx, _, warp_m = mfma_emitter.extract_thread_binding(thread_binding)

                T.clear(C_local)
                T.clear(C_local_accum)

                for ko in T.Pipelined((K // block_K), num_stages=0):

                    # Load A into shared memory
                    T.copy(A[by * block_M, ko * block_K], A_shared)

                    # Load B into shared memory
                    T.copy(B[bx * block_N, ko * block_K], B_shared)


                    for ki in T.serial(0, (block_K // (k_pack * micro_size_k))):

                        # Load A into fragment
                        mfma_emitter.ldmatrix_a(
                            A_local,
                            A_shared,
                            ki,
                        )

                        # Load B into fragment
                        mfma_emitter.ldmatrix_b(
                            B_local,
                            B_shared,
                            ki,
                        )

                        # Perform Matrix Multiplication
                        mfma_emitter.mfma(A_local, B_local, C_local)
                        
                    for wr in T.serial(warp_rows):
                        for wc in T.serial(warp_cols):
                            for local_id in T.serial(local_size_c):
                                i = wr * (warp_cols * local_size_c) + wc * local_size_c + local_id
                                # compute the m idx of the real workloads
                                local_row, _ = thread_id_shared_access_64x4_to_16x16_layout_C_n_m(tx, local_id)
                                m = by * block_M + warp_m * (warp_rows * 16) + wr * 16 + local_row
                                C_local_accum[i] += C_local[i] * A_scale[m, ko] * B_scale[bx, ko]
                    T.clear(C_local)

                # Perform STMatrix
                if cache_write_shared:
                    mfma_emitter.stmatrix(
                        C_local_accum,
                        C_shared,
                    )

                    # Store shared into global
                    for i, j in T.Parallel(block_M, block_N):
                        C[by * block_M + i, bx * block_N + j] = C_shared[
                            i // micro_size_x,
                            j // micro_size_y,
                            i % micro_size_x,
                            j % micro_size_y,
                        ]
                else:
                    mfma_emitter.stmatrix(
                        C_local_accum,
                        C,
                        pid_m=by,
                        pid_n=bx,
                    )

        return main


    autotuner = AutoTuner.from_kernel(
        kernel=kernel, configs=get_configs(M, N, K, with_roller)).set_compile_args(
            out_idx=[-1],
            target="auto",
        ).set_profile_args(
            supply_type=tl.TensorSupplyType.Integer,
            ref_prog=ref_program,
            skip_check=True,
        )
    return autotuner.run(warmup=3, rep=20)

@tilelang.jit
def tl_matmul(
    M,
    N,
    K,
    in_dtype="float8_e4m3fnuz",
    out_dtype="bfloat16",
    accum_dtype="float32",
    a_transposed=False,
    b_transposed=True,
):
    micro_size_x = micro_size_y = micro_size_k = 16

    if in_dtype in {"float8_e4m3fnuz", "int8"}:
        micro_size_k = 32
    
    k_pack = 2
    block_row_warps = 2
    block_col_warps = 2
    warp_row_tiles = 64
    warp_col_tiles = 64
    chunk = 128
    shared_scope = "shared"
    cache_write_shared = False

    block_M = block_row_warps * warp_row_tiles
    block_N = block_col_warps * warp_col_tiles
    block_K = chunk

    scale_n =  (N + block_N - 1) // block_N
    scale_k =  (K + block_K - 1) // block_K

    A_shape = (K, M) if a_transposed else (M, K)
    B_shape = (N, K) if b_transposed else (K, N)
    A_scale_shape = (M, scale_k)
    B_scale_shape = (scale_n, scale_k)

    A_shared_shape = (block_K, block_M) if a_transposed else (block_M, block_K)
    B_shared_shape = (block_N, block_K) if b_transposed else (block_K, block_N)
    C_shared_shape = (
        block_N // micro_size_y,
        micro_size_x,
        micro_size_y,
    )

    warp_size = 64
    threads = warp_size * (block_row_warps * block_col_warps)
    local_size_a = (k_pack * micro_size_x * micro_size_k) // warp_size
    local_size_b = (k_pack * micro_size_y * micro_size_k) // warp_size
    local_size_c = (micro_size_x * micro_size_y) // warp_size
    warp_rows = warp_row_tiles // micro_size_x
    warp_cols = warp_col_tiles // micro_size_y

    # MMA Wrapper to Auto Generate Code for MMA
    mfma_emitter = MatrixCoreIntrinEmitter(
        a_dtype=in_dtype,
        b_dtype=in_dtype,
        accum_dtype=accum_dtype,
        a_transposed=a_transposed,
        b_transposed=b_transposed,
        block_row_warps=block_row_warps,
        block_col_warps=block_col_warps,
        warp_row_tiles=warp_row_tiles,
        warp_col_tiles=warp_col_tiles,
        chunk=chunk,
        k_pack=k_pack,
    )

    @T.prim_func
    def main(
            A: T.Tensor(A_shape, in_dtype),
            B: T.Tensor(B_shape, in_dtype),
            A_scale: T.Tensor(A_scale_shape, accum_dtype),
            B_scale: T.Tensor(B_scale_shape, accum_dtype),
            C: T.Tensor((M, N), out_dtype),
    ):
        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):

            A_shared = T.alloc_shared(A_shared_shape, in_dtype, scope=shared_scope)
            B_shared = T.alloc_shared(B_shared_shape, in_dtype, scope=shared_scope)
            C_shared = T.alloc_shared(C_shared_shape, out_dtype, scope=shared_scope)
            A_local = T.alloc_local((warp_rows * local_size_a), in_dtype)
            B_local = T.alloc_local((warp_cols * local_size_b), in_dtype)
            C_local = T.alloc_local((warp_rows * warp_cols * local_size_c), accum_dtype)
            C_local_accum = T.alloc_local((warp_rows * warp_cols * local_size_c), accum_dtype)

            T.annotate_layout({
                A_shared: make_swizzle_layout(A_shared),
                B_shared: make_swizzle_layout(B_shared),
            })

            # Improve L2 Cache
            T.use_swizzle(panel_size=10)
            
            thread_binding = T.get_thread_binding()
            tx, _, warp_m = mfma_emitter.extract_thread_binding(thread_binding)

            T.clear(C_local)
            T.clear(C_local_accum)

            for ko in T.Pipelined((K // block_K), num_stages=0):

                # Load A into shared memory
                T.copy(A[by * block_M, ko * block_K], A_shared)

                # Load B into shared memory
                T.copy(B[bx * block_N, ko * block_K], B_shared)


                for ki in T.serial(0, (block_K // (k_pack * micro_size_k))):

                    # Load A into fragment
                    mfma_emitter.ldmatrix_a(
                        A_local,
                        A_shared,
                        ki,
                    )

                    # Load B into fragment
                    mfma_emitter.ldmatrix_b(
                        B_local,
                        B_shared,
                        ki,
                    )

                    # Perform Matrix Multiplication
                    mfma_emitter.mfma(A_local, B_local, C_local)
                    
                for wr in T.serial(warp_rows):
                    for wc in T.serial(warp_cols):
                        for local_id in T.serial(local_size_c):
                            i = wr * (warp_cols * local_size_c) + wc * local_size_c + local_id
                            # compute the m idx of the real workloads
                            local_row, _ = thread_id_shared_access_64x4_to_16x16_layout_C_n_m(tx, local_id)
                            m = by * block_M + warp_m * (warp_rows * 16) + wr * 16 + local_row
                            C_local_accum[i] += C_local[i] * A_scale[m, ko] * B_scale[bx, ko]
                T.clear(C_local)

            # Perform STMatrix
            if cache_write_shared:
                mfma_emitter.stmatrix(
                    C_local_accum,
                    C_shared,
                )

                # Store shared into global
                for i, j in T.Parallel(block_M, block_N):
                    C[by * block_M + i, bx * block_N + j] = C_shared[
                        i // micro_size_x,
                        j // micro_size_y,
                        i % micro_size_x,
                        j % micro_size_y,
                    ]
            else:
                mfma_emitter.stmatrix(
                    C_local_accum,
                    C,
                    pid_m=by,
                    pid_n=bx,
                )

    return main

def custom_kernel(data: input_t) -> output_t:
    # (m, k), (n, k), (m, scale_k), (scale_n, scale_k)
    a, b, a_scale, b_scale, c = data
    A = a.contiguous()
    B = b.contiguous()
    scale_A = a_scale.contiguous()
    scale_B = b_scale.contiguous()

    M, N, K = A.shape[0], B.shape[0], A.shape[1]
    in_dtype = "float8_e4m3fnuz"
    out_dtype = "bfloat16"
    accum_dtype = "float32"
    a_transposed = False
    b_transposed = True

    # result = get_best_config(M, N, K, with_roller=False)
    # print(result.config)
    # kernel = result.kernel
    # c = kernel(A, B, scale_A, scale_B)

    kernel = tl_matmul(M, N, K, in_dtype, out_dtype, accum_dtype, a_transposed, b_transposed)
    kernel(A, B, scale_A, scale_B, c)
    
    return c

if __name__ == "__main__":
    data = generate_input(1024, 576, 7168, 42)
    c = custom_kernel(data) 
    print("custom_kernel c: ", c)
    ref = ref_kernel(data)
    print("ref_kernel c: ", ref)
    torch.testing.assert_close(c, ref, rtol=1e-2, atol=1e-2)

python eval.py leaderboard /tmp/tmplo5zasih
WARNING: pip is being invoked by an old script wrapper. This will fail in a future version of pip.
Please see https://github.com/pypa/pip/issues/5599 for advice on fixing the underlying issue.
To avoid this problem you can invoke Python with '-m pip' instead of running pip directly.
WARNING: pip is being invoked by an old script wrapper. This will fail in a future version of pip.
Please see https://github.com/pypa/pip/issues/5599 for advice on fixing the underlying issue.
To avoid this problem you can invoke Python with '-m pip' instead of running pip directly.

Defaulting to user installation because normal site-packages is not writeable
Requirement already satisfied: tilelang-rocm in /home/runner/.local/lib/python3.10/site-packages (0.1.4.post11)
Requirement already satisfied: Cython in /home/runner/.local/lib/python3.10/site-packages (from tilelang-rocm) (3.1.1)
Requirement already satisfied: numpy>=1.23.5 in /usr/local/lib/python3.10/dist-packages (from tilelang-rocm) (1.26.4)
Requirement already satisfied: tqdm>=4.62.3 in /home/runner/.local/lib/python3.10/site-packages (from tilelang-rocm) (4.67.1)
Requirement already satisfied: typing_extensions>=4.10.0 in /usr/local/lib/python3.10/dist-packages (from tilelang-rocm) (4.12.2)
Requirement already satisfied: cloudpickle in /home/runner/.local/lib/python3.10/site-packages (from tilelang-rocm) (3.1.1)
Requirement already satisfied: ml_dtypes in /home/runner/.local/lib/python3.10/site-packages (from tilelang-rocm) (0.5.1)
Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from tilelang-rocm) (7.0.0)
Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (from tilelang-rocm) (2.8.0.dev20250530+rocm6.3)
Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch->tilelang-rocm) (3.16.1)
Requirement already satisfied: sympy>=1.13.3 in /usr/local/lib/python3.10/dist-packages (from torch->tilelang-rocm) (1.13.3)
Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch->tilelang-rocm) (3.4.2)
Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch->tilelang-rocm) (3.1.4)
Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch->tilelang-rocm) (2024.10.0)
Requirement already satisfied: pytorch-triton-rocm==3.3.1+gitc8757738 in /usr/local/lib/python3.10/dist-packages (from torch->tilelang-rocm) (3.3.1+gitc8757738)
Requirement already satisfied: setuptools>=40.8.0 in /usr/lib/python3/dist-packages (from pytorch-triton-rocm==3.3.1+gitc8757738->torch->tilelang-rocm) (59.6.0)
Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy>=1.13.3->torch->tilelang-rocm) (1.3.0)
Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch->tilelang-rocm) (2.1.5)
Defaulting to user installation because normal site-packages is not writeable
Requirement already satisfied: apache-tvm in /home/runner/.local/lib/python3.10/site-packages (0.11.1)
Requirement already satisfied: attrs in /home/runner/.local/lib/python3.10/site-packages (from apache-tvm) (25.3.0)
Requirement already satisfied: cloudpickle in /home/runner/.local/lib/python3.10/site-packages (from apache-tvm) (3.1.1)
Requirement already satisfied: decorator in /home/runner/.local/lib/python3.10/site-packages (from apache-tvm) (5.2.1)
Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from apache-tvm) (1.26.4)
Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from apache-tvm) (7.0.0)
Requirement already satisfied: scipy in /home/runner/.local/lib/python3.10/site-packages (from apache-tvm) (1.15.3)
Requirement already satisfied: synr==0.6.0 in /home/runner/.local/lib/python3.10/site-packages (from apache-tvm) (0.6.0)
Requirement already satisfied: tornado in /home/runner/.local/lib/python3.10/site-packages (from apache-tvm) (6.5.1)
Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from apache-tvm) (4.12.2)
