# 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
import tilelang
print(tilelang.__version__)
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()
# disable std out
import sys
import os
original_stdout = sys.stdout
sys.stdout = open(os.devnull, 'w')

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]
    split_k = [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, split_k, 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],
        "split_k": 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

@tilelang.jit(
    out_idx=-1,
)
def tl_matmul_no_tune(
    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

from tilelang.autotuner import AutoTuner
def get_best_config(M, N, K):

    def kernel(
        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

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

best_configs = {
    (6144, 4096, 7168): {'block_row_warps': 8, 'block_col_warps': 1, 'warp_row_tiles': 16, 'warp_col_tiles': 128, 'split_k': 1, 'num_stages': 1, 'enable_rasterization': False},
    (6144, 7168, 2048): {'block_row_warps': 8, 'block_col_warps': 1, 'warp_row_tiles': 16, 'warp_col_tiles': 128, 'split_k': 1, 'num_stages': 1, 'enable_rasterization': False},
    (6144, 7168, 2304): {'block_row_warps': 4, 'block_col_warps': 2, 'warp_row_tiles': 32, 'warp_col_tiles': 64, 'split_k': 1, 'num_stages': 1, 'enable_rasterization': False},
    (6144, 512, 7168): {'block_row_warps': 4, 'block_col_warps': 1, 'warp_row_tiles': 32, 'warp_col_tiles': 128, 'split_k': 1, 'num_stages': 1, 'enable_rasterization': True},
    (6144, 4096, 512): {'block_row_warps': 2, 'block_col_warps': 2, 'warp_row_tiles': 64, 'warp_col_tiles': 64, 'split_k': 1, 'num_stages': 0, 'enable_rasterization': False},
    (1024, 1536, 7168): {'block_row_warps': 2, 'block_col_warps': 4, 'warp_row_tiles': 16, 'warp_col_tiles': 32, 'split_k': 1, 'num_stages': 1, 'enable_rasterization': False},
    (1024, 3072, 1536): {'block_row_warps': 2, 'block_col_warps': 2, 'warp_row_tiles': 32, 'warp_col_tiles': 64, 'split_k': 1, 'num_stages': 1, 'enable_rasterization': False},
    (1024, 576, 7168): {'block_row_warps': 1, 'block_col_warps': 8, 'warp_row_tiles': 16, 'warp_col_tiles': 16, 'split_k': 1, 'num_stages': 1, 'enable_rasterization': True},
    (1024, 7168, 256): {'block_row_warps': 4, 'block_col_warps': 4, 'warp_row_tiles': 32, 'warp_col_tiles': 32, 'split_k': 1, 'num_stages': 0, 'enable_rasterization': False},
    (1024, 7168, 2048): {'block_row_warps': 8, 'block_col_warps': 1, 'warp_row_tiles': 16, 'warp_col_tiles': 128, 'split_k': 1, 'num_stages': 1, 'enable_rasterization': False},
    (1024, 4608, 7168): {'block_row_warps': 4, 'block_col_warps': 1, 'warp_row_tiles': 32, 'warp_col_tiles': 128, 'split_k': 1, 'num_stages': 1, 'enable_rasterization': False},
    (1024, 7168, 2304): {'block_row_warps': 4, 'block_col_warps': 2, 'warp_row_tiles': 32, 'warp_col_tiles': 64, 'split_k': 1, 'num_stages': 0, 'enable_rasterization': False},
    (1024, 512, 7168): {'block_row_warps': 1, 'block_col_warps': 4, 'warp_row_tiles': 16, 'warp_col_tiles': 32, 'split_k': 1, 'num_stages': 0, 'enable_rasterization': False},
    (1024, 4096, 512): {'block_row_warps': 2, 'block_col_warps': 4, 'warp_row_tiles': 32, 'warp_col_tiles': 32, 'split_k': 1, 'num_stages': 0, 'enable_rasterization': False},
    (6144, 1536, 7168): {'block_row_warps': 4, 'block_col_warps': 2, 'warp_row_tiles': 32, 'warp_col_tiles': 64, 'split_k': 1, 'num_stages': 0, 'enable_rasterization': True},
    (6144, 3072, 1536): {'block_row_warps': 4, 'block_col_warps': 2, 'warp_row_tiles': 32, 'warp_col_tiles': 64, 'split_k': 1, 'num_stages': 0, 'enable_rasterization': False},
    (6144, 576, 7168): {'block_row_warps': 8, 'block_col_warps': 2, 'warp_row_tiles': 16, 'warp_col_tiles': 64, 'split_k': 1, 'num_stages': 0, 'enable_rasterization': True},
    (6144, 7168, 256): {'block_row_warps': 4, 'block_col_warps': 4, 'warp_row_tiles': 32, 'warp_col_tiles': 32, 'split_k': 1, 'num_stages': 0, 'enable_rasterization': False}
}

# from functools import lru_cache
# @lru_cache
def host_func(M, N, K):
    in_dtype = "float8_e4m3fnuz"
    accum_dtype = "float32"
    a_transposed = False
    b_transposed = True
    print(f"M: {M}, N: {N}, K: {K}")
    print(best_configs.get((M, N, K), None))
    if best_configs.get((M, N, K), None) is not None:
        return tl_matmul_no_tune(M, N, K, in_dtype, accum_dtype, accum_dtype, a_transposed, b_transposed, **best_configs[(M, N, K)])
    else:
        return tl_matmul(M, N, K, in_dtype, accum_dtype, accum_dtype, a_transposed, b_transposed)

def custom_kernel(data: input_t) -> output_t:
    A, B, scale_A, scale_B, C = data
    M, N, K = A.shape[0], B.shape[0], A.shape[1]
    print(f"M: {M}, N: {N}, K: {K}")
        
    #  if best config is not None, use it
    kernel = host_func(M, N, K)

    A = A.contiguous()
    B = B.contiguous()
    scale_A = scale_A.contiguous()
    scale_B = scale_B.contiguous()
    
    Split_C = kernel(A, B, scale_A, scale_B, skip_check=True)
    
    C = Split_C
    C = C.sum(dim=0)
    C = C.bfloat16()
    return C
python eval.py benchmark /tmp/tmp7qvszj64
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<21:43,  8.20s/it]
Compiling configurations:  69%|██████▉   | 110/160 [00:08<00:02, 18.76it/s]
Compiling configurations: 100%|██████████| 160/160 [00:09<00:00, 16.93it/s]

Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=1.29]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=1.29]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=1.29]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=1.29]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=1.28]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=1.28]
Bench configurations:   0%|          | 0/160 [00:00<?, ?it/s, best_latency=1.28]
Bench configurations:   4%|▍         | 7/160 [00:00<00:02, 66.51it/s, best_latency=1.28]
Bench configurations:   4%|▍         | 7/160 [00:00<00:02, 66.51it/s, best_latency=1.28]
Bench configurations:   4%|▍         | 7/160 [00:00<00:02, 66.51it/s, best_latency=1.28]
Bench configurations:   4%|▍         | 7/160 [00:00<00:02, 66.51it/s, best_latency=1.28]
Bench configurations:   4%|▍         | 7/160 [00:00<00:02, 66.51it/s, best_latency=1.28]
Bench configurations:   4%|▍         | 7/160 [00:00<00:02, 66.51it/s, best_latency=1.28]
Bench configurations:   4%|▍         | 7/160 [00:00<00:02, 66.51it/s, best_latency=1.28]
Bench configurations:   4%|▍         | 7/160 [00:00<00:02, 66.51it/s, best_latency=1.28]
Bench configurations:   4%|▍         | 7/160 [00:00<00:02, 66.51it/s, best_latency=1.28]
Bench configurations:   4%|▍         | 7/160 [00:00<00:02, 66.51it/s, best_latency=1.28]
Bench configurations:   4%|▍         | 7/160 [00:00<00:02, 66.51it/s, best_latency=1.28]
Bench configurations:  11%|█         | 17/160 [00:00<00:01, 81.42it/s, best_latency=1.28]
Bench configurations:  11%|█         | 17/160 [00:00<00:01, 81.42it/s, best_latency=1.28]
Bench configurations:  11%|█         | 17/160 [00:00<00:01, 81.42it/s, best_latency=1.28]
Bench configurations:  11%|█         | 17/160 [00:00<00:01, 81.42it/s, best_latency=1.28]
Bench configurations:  11%|█         | 17/160 [00:00<00:01, 81.42it/s, best_latency=1.28]
Bench configurations:  11%|█         | 17/160 [00:00<00:01, 81.42it/s, best_latency=1.28]
Bench configurations:  11%|█         | 17/160 [00:00<00:01, 81.42it/s, best_latency=1.28]
Bench configurations:  11%|█         | 17/160 [00:00<00:01, 81.42it/s, best_latency=1.28]
Bench configurations:  11%|█         | 17/160 [00:00<00:01, 81.42it/s, best_latency=1.28]
Bench configurations:  11%|█         | 17/160 [00:00<00:01, 81.42it/s, best_latency=1.28]
Bench configurations:  11%|█         | 17/160 [00:00<00:01, 81.42it/s, best_latency=1.28]
Bench configurations:  17%|█▋        | 27/160 [00:00<00:01, 86.13it/s, best_latency=1.28]
Bench configurations:  17%|█▋        | 27/160 [00:00<00:01, 86.13it/s, best_latency=1.28]
Bench configurations:  17%|█▋        | 27/160 [00:00<00:01, 86.13it/s, best_latency=1.28]
Bench configurations:  17%|█▋        | 27/160 [00:00<00:01, 86.13it/s, best_latency=1.28]
Bench configurations:  17%|█▋        | 27/160 [00:00<00:01, 86.13it/s, best_latency=1.28]
Bench configurations:  17%|█▋        | 27/160 [00:00<00:01, 86.13it/s, best_latency=1.28]
Bench configurations:  17%|█▋        | 27/160 [00:00<00:01, 86.13it/s, best_latency=1.28]
Bench configurations:  17%|█▋        | 27/160 [00:00<00:01, 86.13it/s, best_latency=1.28]
Bench configurations:  17%|█▋        | 27/160 [00:00<00:01, 86.13it/s, best_latency=1.28]
Bench configurations:  17%|█▋        | 27/160 [00:00<00:01, 86.13it/s, best_latency=1.28]
Bench configurations:  17%|█▋        | 27/160 [00:00<00:01, 86.13it/s, best_latency=1.28]
Bench configurations:  23%|██▎       | 37/160 [00:00<00:01, 88.24it/s, best_latency=1.28]
Bench configurations:  23%|██▎       | 37/160 [00:00<00:01, 88.24it/s, best_latency=1.28]
Bench configurations:  23%|██▎       | 37/160 [00:00<00:01, 88.24it/s, best_latency=1.28]
Bench configurations:  23%|██▎       | 37/160 [00:00<00:01, 88.24it/s, best_latency=1.28]
Bench configurations:  23%|██▎       | 37/160 [00:00<00:01, 88.24it/s, best_latency=1.28]
Bench configurations:  23%|██▎       | 37/160 [00:00<00:01, 88.24it/s, best_latency=1.28]
Bench configurations:  23%|██▎       | 37/160 [00:00<00:01, 88.24it/s, best_latency=1.28]
Bench configurations:  23%|██▎       | 37/160 [00:00<00:01, 88.24it/s, best_latency=1.28]
Bench configurations:  23%|██▎       | 37/160 [00:00<00:01, 88.24it/s, best_latency=1.28]
Bench configurations:  23%|██▎       | 37/160 [00:00<00:01, 88.24it/s, best_latency=1.28]
Bench configurations:  23%|██▎       | 37/160 [00:00<00:01, 88.24it/s, best_latency=1.28]
Bench configurations:  29%|██▉       | 47/160 [00:00<00:01, 89.46it/s, best_latency=1.28]
Bench configurations:  29%|██▉       | 47/160 [00:00<00:01, 89.46it/s, best_latency=1.28]
Bench configurations:  29%|██▉       | 47/160 [00:00<00:01, 89.46it/s, best_latency=1.28]
Bench configurations:  29%|██▉       | 47/160 [00:00<00:01, 89.46it/s, best_latency=1.28]
Bench configurations:  29%|██▉       | 47/160 [00:00<00:01, 89.46it/s, best_latency=1.28]
Bench configurations:  29%|██▉       | 47/160 [00:00<00:01, 89.46it/s, best_latency=1.28]
Bench configurations:  29%|██▉       | 47/160 [00:00<00:01, 89.46it/s, best_latency=1.28]
Bench configurations:  29%|██▉       | 47/160 [00:00<00:01, 89.46it/s, best_latency=1.28]
Bench configurations:  29%|██▉       | 47/160 [00:00<00:01, 89.46it/s, best_latency=1.28]
Bench configurations:  29%|██▉       | 47/160 [00:00<00:01, 89.46it/s, best_latency=1.28]
Bench configurations:  29%|██▉       | 47/160 [00:00<00:01, 89.46it/s, best_latency=1.28]
Bench configurations:  36%|███▌      | 57/160 [00:00<00:01, 90.16it/s, best_latency=1.28]
Bench configurations:  36%|███▌      | 57/160 [00:00<00:01, 90.16it/s, best_latency=1.28]
Bench configurations:  36%|███▌      | 57/160 [00:00<00:01, 90.16it/s, best_latency=1.28]
Bench configurations:  36%|███▌      | 57/160 [00:00<00:01, 90.16it/s, best_latency=1.28]
Bench configurations:  36%|███▌      | 57/160 [00:00<00:01, 90.16it/s, best_latency=1.28]
Bench configurations:  36%|███▌      | 57/160 [00:00<00:01, 90.16it/s, best_latency=1.28]
Bench configurations:  36%|███▌      | 57/160 [00:00<00:01, 90.16it/s, best_latency=1.28]
Bench configurations:  36%|███▌      | 57/160 [00:00<00:01, 90.16it/s, best_latency=1.28]
Bench configurations:  36%|███▌      | 57/160 [00:00<00:01, 90.16it/s, best_latency=1.28]
Bench configurations:  36%|███▌      | 57/160 [00:00<00:01, 90.16it/s, best_latency=1.28]
Bench configurations:  36%|███▌      | 57/160 [00:00<00:01, 90.16it/s, best_latency=1.28]
Bench configurations:  42%|████▏     | 67/160 [00:00<00:01, 90.64it/s, best_latency=1.28]
Bench configurations:  42%|████▏     | 67/160 [00:00<00:01, 90.64it/s, best_latency=1.28]
Bench configurations:  42%|████▏     | 67/160 [00:00<00:01, 90.64it/s, best_latency=1.28]
Bench configurations:  42%|████▏     | 67/160 [00:00<00:01, 90.64it/s, best_latency=1.28]
Bench configurations:  42%|████▏     | 67/160 [00:00<00:01, 90.64it/s, best_latency=1.28]
Bench configurations:  42%|████▏     | 67/160 [00:00<00:01, 90.64it/s, best_latency=1.28]
Bench configurations:  42%|████▏     | 67/160 [00:00<00:01, 90.64it/s, best_latency=1.28]
Bench configurations:  42%|████▏     | 67/160 [00:00<00:01, 90.64it/s, best_latency=1.28]
Bench configurations:  42%|████▏     | 67/160 [00:00<00:01, 90.64it/s, best_latency=1.28]
Bench configurations:  42%|████▏     | 67/160 [00:00<00:01, 90.64it/s, best_latency=1.28]
Bench configurations:  42%|████▏     | 67/160 [00:00<00:01, 90.64it/s, best_latency=1.28]
Bench configurations:  48%|████▊     | 77/160 [00:00<00:00, 90.94it/s, best_latency=1.28]
Bench configurations:  48%|████▊     | 77/160 [00:00<00:00, 90.94it/s, best_latency=1.28]
Bench configurations:  48%|████▊     | 77/160 [00:00<00:00, 90.94it/s, best_latency=1.28]
Bench configurations:  48%|████▊     | 77/160 [00:00<00:00, 90.94it/s, best_latency=1.28]
Bench configurations:  48%|████▊     | 77/160 [00:00<00:00, 90.94it/s, best_latency=1.28]
Bench configurations:  48%|████▊     | 77/160 [00:00<00:00, 90.94it/s, best_latency=1.28]
Bench configurations:  48%|████▊     | 77/160 [00:00<00:00, 90.94it/s, best_latency=1.28]
Bench configurations:  48%|████▊     | 77/160 [00:00<00:00, 90.94it/s, best_latency=1.28]
Bench configurations:  48%|████▊     | 77/160 [00:00<00:00, 90.94it/s, best_latency=1.28]
Bench configurations:  48%|████▊     | 77/160 [00:00<00:00, 90.94it/s, best_latency=1.28]
Bench configurations:  48%|████▊     | 77/160 [00:00<00:00, 90.94it/s, best_latency=1.28]
Bench configurations:  54%|█████▍    | 87/160 [00:00<00:00, 91.13it/s, best_latency=1.28]
Bench configurations:  54%|█████▍    | 87/160 [00:00<00:00, 91.13it/s, best_latency=1.28]
Bench configurations:  54%|█████▍    | 87/160 [00:01<00:00, 91.13it/s, best_latency=1.28]
Bench configurations:  54%|█████▍    | 87/160 [00:01<00:00, 91.13it/s, best_latency=1.28]
Bench configurations:  54%|█████▍    | 87/160 [00:01<00:00, 91.13it/s, best_latency=1.28]
Bench configurations:  54%|█████▍    | 87/160 [00:01<00:00, 91.13it/s, best_latency=1.28]
Bench configurations:  54%|█████▍    | 87/160 [00:01<00:00, 91.13it/s, best_latency=1.28]
Bench configurations:  54%|█████▍    | 87/160 [00:01<00:00, 91.13it/s, best_latency=1.28]
Bench configurations:  54%|█████▍    | 87/160 [00:01<00:00, 91.13it/s, best_latency=1.28]
Bench configurations:  54%|█████▍    | 87/160 [00:01<00:00, 91.13it/s, best_latency=1.28]
Bench configurations:  54%|█████▍    | 87/160 [00:01<00:00, 91.13it/s, best_latency=1.28]
Bench configurations:  61%|██████    | 97/160 [00:01<00:00, 91.27it/s, best_latency=1.28]
Bench configurations:  61%|██████    | 97/160 [00:01<00:00, 91.27it/s, best_latency=1.28]
Bench configurations:  61%|██████    | 97/160 [00:01<00:00, 91.27it/s, best_latency=1.28]
Bench configurations:  61%|██████    | 97/160 [00:01<00:00, 91.27it/s, best_latency=1.28]
Bench configurations:  61%|██████    | 97/160 [00:01<00:00, 91.27it/s, best_latency=1.28]
Bench configurations:  61%|██████    | 97/160 [00:01<00:00, 91.27it/s, best_latency=1.28]
Bench configurations:  61%|██████    | 97/160 [00:01<00:00, 91.27it/s, best_latency=1.28]
Bench configurations:  61%|██████    | 97/160 [00:01<00:00, 91.27it/s, best_latency=1.28]
Bench configurations:  61%|██████    | 97/160 [00:01<00:00, 91.27it/s, best_latency=1.28]
Bench configurations:  61%|██████    | 97/160 [00:01<00:00, 91.27it/s, best_latency=1.28]
Bench configurations:  61%|██████    | 97/160 [00:01<00:00, 91.27it/s, best_latency=1.2] 
Bench configurations:  67%|██████▋   | 107/160 [00:01<00:00, 86.53it/s, best_latency=1.2]
Bench configurations:  67%|██████▋   | 107/160 [00:01<00:00, 86.53it/s, best_latency=1.2]
Bench configurations:  67%|██████▋   | 107/160 [00:01<00:00, 86.53it/s, best_latency=1.2]
Bench configurations:  67%|██████▋   | 107/160 [00:01<00:00, 86.53it/s, best_latency=1.2]
Bench configurations:  67%|██████▋   | 107/160 [00:01<00:00, 86.53it/s, best_latency=1.2]
Bench configurations:  67%|██████▋   | 107/160 [00:01<00:00, 86.53it/s, best_latency=0.882]
Bench configurations:  67%|██████▋   | 107/160 [00:01<00:00, 86.53it/s, best_latency=0.882]
Bench configurations:  67%|██████▋   | 107/160 [00:01<00:00, 86.53it/s, best_latency=0.882]
Bench configurations:  67%|██████▋   | 107/160 [00:01<00:00, 86.53it/s, best_latency=0.882]
Bench configurations:  67%|██████▋   | 107/160 [00:01<00:00, 86.53it/s, best_latency=0.882]
Bench configurations:  72%|███████▎  | 116/160 [00:01<00:00, 83.27it/s, best_latency=0.882]
Bench configurations:  72%|███████▎  | 116/160 [00:01<00:00, 83.27it/s, best_latency=0.882]
Bench configurations:  72%|███████▎  | 116/160 [00:01<00:00, 83.27it/s, best_latency=0.882]
Bench configurations:  72%|███████▎  | 116/160 [00:01<00:00, 83.27it/s, best_latency=0.882]
Bench configurations:  72%|███████▎  | 116/160 [00:01<00:00, 83.27it/s, best_latency=0.882]
Bench configurations:  72%|███████▎  | 116/160 [00:01<00:00, 83.27it/s, best_latency=0.882]
Bench configurations:  72%|███████▎  | 116/160 [00:01<00:00, 83.27it/s, best_latency=0.882]
Bench configurations:  72%|███████▎  | 116/160 [00:01<00:00, 83.27it/s, best_latency=0.882]
Bench configurations:  72%|███████▎  | 116/160 [00:01<00:00, 83.27it/s, best_latency=0.882]
Bench configurations:  72%|███████▎  | 116/160 [00:01<00:00, 83.27it/s, best_latency=0.882]
Bench configurations:  78%|███████▊  | 125/160 [00:01<00:00, 75.38it/s, best_latency=0.882]
Bench configurations:  78%|███████▊  | 125/160 [00:01<00:00, 75.38it/s, best_latency=0.882]
Bench configurations:  78%|███████▊  | 125/160 [00:01<00:00, 75.38it/s, best_latency=0.882]
Bench configurations:  78%|███████▊  | 125/160 [00:01<00:00, 75.38it/s, best_latency=0.882]
Bench configurations:  78%|███████▊  | 125/160 [00:01<00:00, 75.38it/s, best_latency=0.882]
Bench configurations:  78%|███████▊  | 125/160 [00:01<00:00, 75.38it/s, best_latency=0.882]
Bench configurations:  78%|███████▊  | 125/160 [00:01<00:00, 75.38it/s, best_latency=0.882]
Bench configurations:  78%|███████▊  | 125/160 [00:01<00:00, 75.38it/s, best_latency=0.882]
Bench configurations:  78%|███████▊  | 125/160 [00:01<00:00, 75.38it/s, best_latency=0.882]
Bench configurations:  83%|████████▎ | 133/160 [00:01<00:00, 74.18it/s, best_latency=0.882]
Bench configurations:  83%|████████▎ | 133/160 [00:01<00:00, 74.18it/s, best_latency=0.882]
Bench configurations:  83%|████████▎ | 133/160 [00:01<00:00, 74.18it/s, best_latency=0.882]
Bench configurations:  83%|████████▎ | 133/160 [00:01<00:00, 74.18it/s, best_latency=0.882]
Bench configurations:  83%|████████▎ | 133/160 [00:01<00:00, 74.18it/s, best_latency=0.882]
Bench configurations:  83%|████████▎ | 133/160 [00:01<00:00, 74.18it/s, best_latency=0.882]
Bench configurations:  83%|████████▎ | 133/160 [00:01<00:00, 74.18it/s, best_latency=0.882]
Bench configurations:  83%|████████▎ | 133/160 [00:01<00:00, 74.18it/s, best_latency=0.882]
Bench configurations:  83%|████████▎ | 133/160 [00:01<00:00, 74.18it/s, best_latency=0.882]
Bench configurations:  83%|████████▎ | 133/160 [00:01<00:00, 74.18it/s, best_latency=0.882]
Bench configurations:  89%|████████▉ | 142/160 [00:01<00:00, 76.87it/s, best_latency=0.882]
Bench configurations:  89%|████████▉ | 142/160 [00:01<00:00, 76.87it/s, best_latency=0.882]
Bench configurations:  89%|████████▉ | 142/160 [00:01<00:00, 76.87it/s, best_latency=0.882]
Bench configurations:  89%|████████▉ | 142/160 [00:01<00:00, 76.87it/s, best_latency=0.882]
Bench configurations:  89%|████████▉ | 142/160 [00:01<00:00, 76.87it/s, best_latency=0.882]
Bench configurations:  89%|████████▉ | 142/160 [00:01<00:00, 76.87it/s, best_latency=0.882]
Bench configurations:  89%|████████▉ | 142/160 [00:01<00:00, 76.87it/s, best_latency=0.882]
Bench configurations:  89%|████████▉ | 142/160 [00:01<00:00, 76.87it/s, best_latency=0.882]
Bench configurations:  89%|████████▉ | 142/160 [00:01<00:00, 76.87it/s, best_latency=0.882]
Bench configurations:  89%|████████▉ | 142/160 [00:01<00:00, 76.87it/s, best_latency=0.882]
Bench configurations:  94%|█████████▍| 151/160 [00:01<00:00, 79.89it/s, best_latency=0.882]
Bench configurations:  94%|█████████▍| 151/160 [00:01<00:00, 79.89it/s, best_latency=0.882]
Bench configurations:  94%|█████████▍| 151/160 [00:01<00:00, 79.89it/s, best_latency=0.882]
Bench configurations:  94%|█████████▍| 151/160 [00:01<00:00, 79.89it/s, best_latency=0.882]
[...] 9 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)
