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

from pip._internal import main as pip_main
pip_main(["install","-U","tilelang-rocm"])
pip_main(["install","-U","apache-tvm"])
from task import input_t, output_t
from reference import generate_input, ref_kernel
import torch
import tilelang
from tilelang import tvm as tvm
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

tilelang.disable_cache()

from tqdm import tqdm
tqdm.write = lambda s: None 

def get_configs():
    import itertools
    block_row_warps = [1, 2, 4, 8, 16]
    block_col_warps = [1, 2, 4, 8, 16]
    warp_row_tiles = [16, 32, 64, 128]
    warp_col_tiles = [16, 32, 64, 128]
    num_split = [1]
    num_stages = [0, 1]
    enable_rasterization = [False, True]

    _configs = list(itertools.product(block_row_warps, block_col_warps, warp_row_tiles, warp_col_tiles, num_split, num_stages, enable_rasterization))
    # filter configs
    _configs = [c for c in _configs if c[0] * c[2] <= 128]
    _configs = [c for c in _configs if c[1] * c[3] == 128]
    
    return [{
        "block_row_warps": c[0],
        "block_col_warps": c[1],
        "warp_row_tiles": c[2],
        "warp_col_tiles": c[3],
        "num_split": c[4],
        "num_stages": c[5],
        "enable_rasterization": c[6],
    } for c in _configs]

@tilelang.autotune(
    configs=get_configs(),
)
@tilelang.jit(
    out_idx=-1,
)
def tl_matmul(
    M,
    N,
    K,
    in_dtype="float8_e4m3fnuz",
    out_dtype="bfloat16",
    accum_dtype="float32",
    a_transposed=False,
    b_transposed=True,
    block_row_warps=2,
    block_col_warps=2,
    warp_row_tiles=64,
    warp_col_tiles=64,
    split_k=2,
    num_stages=1,
    enable_rasterization=False,
):
    micro_size_x = micro_size_y = micro_size_k = 16

    if in_dtype in {"float8_e4m3fnuz", "int8"}:
        micro_size_k = 32
    
    k_pack = 2
    chunk = 128
    shared_scope = "shared.dyn"
    splitK = K // split_k

    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_M // micro_size_x,
        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((split_k, M, N), out_dtype),
    ):
        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), split_k, threads=threads) as (bx, by, bz):

            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(T.ceildiv(splitK, block_K), num_stages=num_stages):

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

                # Load B into shared memory
                T.copy(B[bx * block_N, bz * splitK + 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, bz * T.ceildiv(splitK, block_K)+ ko] * B_scale[bx, bz * T.ceildiv(splitK, block_K)+ ko]

                T.clear(C_local)

            # Perform STMatrix
            mfma_emitter.stmatrix(
                C_local_accum,
                C_shared,
            )

            # Store shared into global
            for i, j in T.Parallel(block_M, block_N):
                C[bz, 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,
                ]

    return main

def custom_kernel(data: input_t) -> output_t:
    A, B, scale_A, scale_B, C = data
    in_dtype = "float8_e4m3fnuz"
    accum_dtype = "float32"
    a_transposed = False
    b_transposed = True
    M, N, K = A.shape[0], B.shape[0], A.shape[1]
    kernel = tl_matmul(M, N, K, in_dtype, accum_dtype, accum_dtype, a_transposed, b_transposed, split_k=2)
    src_code = kernel.get_kernel_source()
    # src_code is the generated cuda source
    assert src_code is not None

    # (m, k), (n, k), (m, scale_k), (scale_n, scale_k)
    def pre_process(A, B, scale_A, scale_B):
        A = A.contiguous()
        B = B.contiguous()
        scale_A = scale_A.contiguous()
        scale_B = scale_B.contiguous()
        return A, B, scale_A, scale_B
    
    A, B, scale_A, scale_B = pre_process(A, B, scale_A, scale_B)
    
    Split_C = kernel(A, B, scale_A, scale_B, skip_check=True)

    def post_process(C):
        C = C.sum(dim=0)
        C = C.bfloat16()
        return C
    
    tilelang_c = post_process(Split_C)
    return tilelang_c
python eval.py test /tmp/tmp977iq_8s
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.

Compiling configurations:   0%|          | 0/160 [00:00<?, ?it/s]
Compiling configurations:   1%|          | 1/160 [00:08<23:25,  8.84s/it]
Compiling configurations:  69%|██████▉   | 111/160 [00:08<00:02, 17.60it/s]
Compiling configurations: 100%|██████████| 160/160 [00:10<00:00, 15.75it/s]

Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00658]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00658]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00633]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00633]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00633]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00633]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00633]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00633]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00613]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00613]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00613]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00613]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00613]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations:  77%|███████▋  | 123/160 [00:00<00:00, 1219.28it/s, best_latency=0.00609]
Bench configurations: 100%|██████████| 160/160 [00:00<00:00, 1104.61it/s, best_latency=0.00609]

Compiling configurations:   0%|          | 0/160 [00:00<?, ?it/s]
Compiling configurations:   1%|          | 1/160 [00:08<21:34,  8.14s/it]
Compiling configurations:  66%|██████▋   | 106/160 [00:08<00:02, 18.09it/s]
Compiling configurations:  93%|█████████▎| 149/160 [00:08<00:00, 26.06it/s]
Compiling configurations: 100%|██████████| 160/160 [00:09<00:00, 17.35it/s]

Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.055]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.055]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.055]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.055]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.0486]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.0486]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.0485]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.0485]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.0485]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=0.0485]
[...] 1684 lines omitted
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.post12)
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.dev20250601+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)
2025-06-02 09:24:19  [TileLang:tilelang.jit.adapter.cython.adapter:INFO]: No cached version found for cython jit adapter, need to compile...
2025-06-02 09:24:19  [TileLang:tilelang.jit.adapter.cython.adapter:INFO]: Waiting for lock to compile cython jit adapter...
2025-06-02 09:24:19  [TileLang:tilelang.jit.adapter.cython.adapter:INFO]: Compiling cython jit adapter...
