TheRock/rocm-libraries/shared/rocroller/scripts/lib/rrperf/rrsuites.py

# Copyright Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT

from itertools import product
from pathlib import Path
from typing import Generator

from rrperf.problems import (
    CodeGenRun,
    GEMMRun,
    MKNLTuple,
    TypeParameters,
)
from rrperf.utils import rocm_gfx

repo_dir = Path(__file__).resolve().parent.parent.parent.parent

fp4fp4_fp32 = TypeParameters(
    type_A="fp4",
    type_B="fp4",
    type_C="float",
    type_D="float",
)

fp6fp6_fp32 = TypeParameters(
    type_A="fp6",
    type_B="fp6",
    type_C="float",
    type_D="float",
)

bf6bf6_fp32 = TypeParameters(
    type_A="bf6",
    type_B="bf6",
    type_C="float",
    type_D="float",
)

fp8fp8_fp32 = TypeParameters(
    type_A="fp8",
    type_B="fp8",
    type_C="float",
    type_D="float",
)

bf8bf8_fp32 = TypeParameters(
    type_A="bf8",
    type_B="bf8",
    type_C="float",
    type_D="float",
)

fp8bf8_fp32 = TypeParameters(
    type_A="fp8",
    type_B="bf8",
    type_C="float",
    type_D="float",
)

bf8fp8_fp32 = TypeParameters(
    type_A="bf8",
    type_B="fp8",
    type_C="float",
    type_D="float",
)

fp16 = TypeParameters(
    type_A="half",
    type_B="half",
    type_C="half",
    type_D="half",
)

bf16_fp32 = TypeParameters(
    type_A="bf16",
    type_B="bf16",
    type_C="float",
    type_D="float",
)

bf16_bf16 = TypeParameters(
    type_A="bf16",
    type_B="bf16",
    type_C="bf16",
    type_D="bf16",
)

fp32 = TypeParameters(
    type_A="float",
    type_B="float",
    type_C="float",
    type_D="float",
)

SGEMM_3072x4096x4096 = dict(
    M=3072, N=4096, K=4096, mac_m=64, mac_n=64, mac_k=64, types=fp32
)

HGEMM_7680x8448x8192 = dict(
    M=7680, N=8448, K=8192, mac_m=64, mac_n=64, mac_k=64, types=fp16
)

HGEMM_7680x8448x8448 = dict(
    M=7680,
    N=8448,
    K=8448,  # matches the guidepost unit test
    mac_m=64,
    mac_n=64,
    mac_k=64,
    workgroup_size_x=128,
    workgroup_size_y=2,
    types=TypeParameters(fp16, trans_A="N", trans_B="T"),
)

SGEMM_256x256x16384 = dict(
    M=256, N=256, K=16384, mac_m=64, mac_n=64, mac_k=64, types=fp32
)


def update_parameters(*args, **kwargs):
    rv = {}
    for d in args:
        rv.update(d)
    rv.update(kwargs)
    return rv


def mkGEMM(*args, **kwargs):
    return GEMMRun(**update_parameters(*args, **kwargs))


def unit():
    default = dict(
        M=1024,
        N=1024,
        K=128,
        mac_m=64,
        mac_n=64,
        mac_k=64,
        numWarmUp=1,
        numOuter=1,
        numInner=1,
    )
    yield mkGEMM(default, types=fp32)
    yield mkGEMM(default, types=fp16)
    yield from tail_loop_reproducer()


def unit_gfx120X():
    default = dict(
        M=1024,
        N=1024,
        K=128,
        mac_m=64,
        mac_n=64,
        mac_k=64,
        wave_m=16,
        wave_n=16,
        wave_k=16,
        numWarmUp=1,
        numOuter=1,
        numInner=1,
    )
    yield mkGEMM(default, types=fp16)
    yield mkGEMM(default, types=bf16_fp32)


def sgemm():
    yield mkGEMM(SGEMM_3072x4096x4096)
    yield mkGEMM(SGEMM_3072x4096x4096, mac_m=128, mac_n=64, mac_k=16)


def hgemm_tensile_guidepost():
    yield mkGEMM(HGEMM_7680x8448x8448)


def hgemm():
    yield mkGEMM(HGEMM_7680x8448x8192)
    yield mkGEMM(
        HGEMM_7680x8448x8192,
        mac_m=128,
        mac_n=256,
        mac_k=16,
        workgroup_size_x=128,
        workgroup_size_y=2,
    )
    yield mkGEMM(
        HGEMM_7680x8448x8192,
        mac_m=128,
        mac_n=256,
        mac_k=16,
        workgroup_size_x=128,
        workgroup_size_y=2,
        prefetchInFlight=2,
        prefetchLDSFactor=2,
        types=TypeParameters(
            HGEMM_7680x8448x8192["types"],
            trans_A="N",
            trans_B="T",
        ),
    )
    yield mkGEMM(
        HGEMM_7680x8448x8192,
        mac_m=128,
        mac_n=256,
        mac_k=16,
        workgroup_size_x=64,
        workgroup_size_y=4,
    )

    yield mkGEMM(
        HGEMM_7680x8448x8192,
        types=TypeParameters(HGEMM_7680x8448x8192["types"], trans_A="T", trans_B="N"),
    )
    yield mkGEMM(
        HGEMM_7680x8448x8192,
        types=TypeParameters(HGEMM_7680x8448x8192["types"], trans_A="T", trans_B="T"),
    )
    yield mkGEMM(
        HGEMM_7680x8448x8192,
        types=TypeParameters(HGEMM_7680x8448x8192["types"], trans_A="N", trans_B="T"),
    )

    yield mkGEMM(HGEMM_7680x8448x8448)
    yield mkGEMM(
        HGEMM_7680x8448x8448,
        mac_m=128,
        mac_n=256,
        mac_k=16,
        workgroup_size_x=256,
        workgroup_size_y=1,
    )
    yield mkGEMM(
        HGEMM_7680x8448x8448,
        mac_m=128,
        mac_n=256,
        mac_k=16,
        workgroup_size_x=128,
        workgroup_size_y=2,
    )
    yield mkGEMM(
        HGEMM_7680x8448x8448,
        mac_m=128,
        mac_n=256,
        mac_k=16,
        workgroup_size_x=64,
        workgroup_size_y=4,
    )
    yield mkGEMM(
        HGEMM_7680x8448x8448,
        mac_m=128,
        mac_n=256,
        mac_k=16,
        workgroup_size_x=128,
        workgroup_size_y=2,
        prefetchInFlight=2,
        prefetchLDSFactor=2,
    )

    for sched in ["Priority", "Cooperative", "Sequential"]:
        yield mkGEMM(
            M=3840,
            N=4224,
            K=4224,
            mac_m=128,
            mac_n=128,
            mac_k=32,
            workgroup_size_x=128,
            workgroup_size_y=2,
            visualize=False,
            scheduler=sched,
            types=TypeParameters(
                fp16,
                trans_A="N",
                trans_B="T",
            ),
        )

        yield mkGEMM(
            M=1024,
            N=50304,
            K=8192,
            mac_m=128,
            mac_n=128,
            mac_k=32,
            workgroup_size_x=128,
            workgroup_size_y=2,
            visualize=False,
            scheduler=sched,
            types=TypeParameters(
                fp16,
                trans_A="N",
                trans_B="T",
            ),
        )

        yield mkGEMM(
            M=3840,
            N=4224,
            K=4224,
            mac_m=128,
            mac_n=128,
            mac_k=32,
            workgroup_size_x=128,
            workgroup_size_y=2,
            betaInFma=False,
            visualize=False,
            scheduler=sched,
            types=TypeParameters(
                fp16,
                trans_A="N",
                trans_B="T",
            ),
        )

        yield mkGEMM(
            M=1024,
            N=50304,
            K=8192,
            mac_m=128,
            mac_n=128,
            mac_k=32,
            workgroup_size_x=128,
            workgroup_size_y=2,
            betaInFma=False,
            visualize=False,
            scheduler=sched,
            types=TypeParameters(
                fp16,
                trans_A="N",
                trans_B="T",
            ),
        )

        yield mkGEMM(
            M=8448,
            N=8448,
            K=128,
            mac_m=128,
            mac_n=256,
            mac_k=16,
            workgroup_size_x=128,
            workgroup_size_y=2,
            visualize=False,
            scheduler=sched,
            types=TypeParameters(
                fp16,
                trans_A="N",
                trans_B="T",
            ),
        )

        yield mkGEMM(
            M=7680,
            N=8448,
            K=8192,
            mac_m=128,
            mac_n=256,
            mac_k=16,
            workgroup_size_x=256,
            workgroup_size_y=1,
            visualize=False,
            scheduler=sched,
            types=TypeParameters(
                fp16,
                trans_A="N",
                trans_B="T",
            ),
        )

    # TODO: Enable once visualizer is working
    # yield from visualizer()


def hgemm_gfx120X():
    params = dict(
        wave_m=16,
        wave_n=16,
        wave_k=16,
        workgroup_size_x=64,
        prefetch=False,
    )

    type_specifiers = [
        ("half", fp16),
        ("bf16", bf16_bf16),
        ("float", fp16),
        ("float", fp8fp8_fp32),
        ("float", fp8bf8_fp32),
        ("float", bf8bf8_fp32),
        ("float", bf8fp8_fp32),
    ]

    for sched in ["Priority", "Cooperative", "Sequential"]:
        for a, b in product("NT", repeat=2):
            for acc, abcd in type_specifiers:
                yield mkGEMM(
                    HGEMM_7680x8448x8192,
                    types=TypeParameters(abcd, type_acc=acc, trans_A=a, trans_B=b),
                    scheduler=sched,
                    **params,
                )


def visualizer():
    yield mkGEMM(
        M=512,
        N=768,
        K=512,
        mac_m=128,
        mac_n=256,
        mac_k=16,
        alpha=2.0,
        beta=0.5,
        workgroup_size_x=256,
        workgroup_size_y=1,
        store="VGPRToGlobalMemoryWithBuffer",
        visualize=True,
        prefetch=False,
        types=TypeParameters(
            fp16,
            trans_A="N",
            trans_B="T",
        ),
    )


def tail_loop_reproducer():
    yield mkGEMM(
        M=64,
        N=128,
        K=8,
        wave_m=32,
        wave_n=32,
        wave_k=2,
        wave_b=1,
        mac_m=64,
        mac_n=64,
        mac_k=8,
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
        ),
    )


def guidepost_1():
    yield mkGEMM(
        HGEMM_7680x8448x8448,
        mac_m=128,
        mac_n=256,
        mac_k=16,
        scheduler="Priority",
        prefetch=True,
        prefetchInFlight=2,
        prefetchLDSFactor=2,
    )


def guidepost_2():
    yield mkGEMM(
        HGEMM_7680x8448x8448,
        K=8192,
        mac_m=128,
        mac_n=256,
        mac_k=16,
        scheduler="Priority",
        prefetch=True,
        prefetchInFlight=2,
        prefetchLDSFactor=2,
    )


def hgemm_no_store_LDS():
    for K in [8448, 8192]:
        yield mkGEMM(
            HGEMM_7680x8448x8448,
            K=K,
            mac_m=128,
            mac_n=256,
            mac_k=16,
            store="VGPRToGlobalMemoryWithBuffer",
            scheduler="Priority",
            prefetch=True,
            prefetchInFlight=2,
            prefetchLDSFactor=0,
        )


def streamk_sweep():
    # Prefetch configurations: (prefetch, prefetchInFlight, prefetchLDSFactor)
    prefetchConfigs = [(False, 0, 0)] + [(True, 2, 2)]

    for twoTile, twoTileDPFirst in [(True, False), (False, True), (False, False)]:
        for base in [HGEMM_7680x8448x8448]:
            # Currently these run out of LDS everywhere except gfx950.
            # + [SGEMM_3072x4096x4096]
            for mac_m in [64, 128]:
                for mac_n in [64, 128, 256]:
                    for mac_k in [16, 32, 64]:
                        for (
                            prefetch,
                            prefetchInFlight,
                            prefetchLDSFactor,
                        ) in prefetchConfigs:
                            # Runs out of VGPRs
                            if (
                                twoTile or twoTileDPFirst
                            ) and mac_m * mac_n * mac_k >= (64 * 256 * 64):
                                continue

                            # Runs out of VGPRs: TwoTile/TwoTileDPFirst + prefetch
                            if (twoTile or twoTileDPFirst) and prefetch:
                                continue

                            # Runs out of VGPRs: Standard + large tiles + prefetchLDSFactor=2
                            if (
                                (not twoTile and not twoTileDPFirst)
                                and prefetch
                                and mac_n == 256
                                and mac_k == 64
                                and prefetchLDSFactor == 2
                            ):
                                continue

                            yield mkGEMM(
                                base,
                                mac_m=mac_m,
                                mac_n=mac_n,
                                mac_k=mac_k,
                                workgroup_size_x=128,
                                workgroup_size_y=2,
                                visualize=False,
                                prefetch=prefetch,
                                prefetchInFlight=prefetchInFlight,
                                prefetchLDSFactor=prefetchLDSFactor,
                                streamK=(
                                    "TwoTile"
                                    if twoTile
                                    else (
                                        "TwoTileDPFirst"
                                        if twoTileDPFirst
                                        else "Standard"
                                    )
                                ),
                                types=TypeParameters(
                                    base["types"],
                                    trans_A="N",
                                    trans_B="T",
                                ),
                            )


def streamk():
    common_overrides = dict(
        workgroup_size_x=128,
        workgroup_size_y=2,
        prefetch=False,
        # TODO: Consider enabling, some run out of VGPRs
        # prefetch=True,
        # prefetchInFlight=2,
        # prefetchLDSFactor=2,
    )

    for streamKMode in ["TwoTile", "TwoTileDPFirst", "Standard"]:
        # SGEMM
        yield mkGEMM(
            SGEMM_3072x4096x4096,
            **common_overrides,
            streamK=streamKMode,
            types=TypeParameters(
                SGEMM_3072x4096x4096["types"],
                trans_A="N",
                trans_B="T",
            ),
        )
        # HGEMM
        yield mkGEMM(
            HGEMM_7680x8448x8448,
            mac_m=128,
            mac_n=256,
            mac_k=16,
            **common_overrides,
            streamK=streamKMode,
            types=TypeParameters(
                HGEMM_7680x8448x8448["types"],
                trans_A="N",
                trans_B="T",
            ),
        )
        yield mkGEMM(
            HGEMM_7680x8448x8192,
            mac_m=128,
            mac_n=256,
            mac_k=16,
            **common_overrides,
            streamK=streamKMode,
            types=TypeParameters(
                HGEMM_7680x8448x8192["types"],
                trans_A="N",
                trans_B="T",
            ),
        )


def smallMN_largeK_fp32():
    yield mkGEMM(
        SGEMM_256x256x16384,
        workgroup_size_x=128,
        workgroup_size_y=2,
        visualize=False,
        prefetch=False,
        # prefetchInFlight=2,
        # prefetchLDSFactor=2,
        streamK="None",
        types=TypeParameters(
            SGEMM_256x256x16384["types"],
            trans_A="T",
            trans_B="N",
        ),
    )


def scalar_is_zero():
    # TODO: Make streamK and ConstantPropagation transformation can be both applied
    sgemm = update_parameters(
        SGEMM_3072x4096x4096,
        beta=0.0,
        streamK="None",
    )
    yield mkGEMM(sgemm)
    yield mkGEMM(sgemm, mac_m=128, mac_n=64, mac_k=16)

    hgemm = update_parameters(
        HGEMM_7680x8448x8192,
        beta=0.0,
        streamK="None",
    )
    yield mkGEMM(hgemm)
    yield mkGEMM(
        hgemm,
        mac_m=128,
        mac_n=256,
        mac_k=16,
        workgroup_size_x=128,
        workgroup_size_y=2,
    )
    yield mkGEMM(
        hgemm,
        mac_m=128,
        mac_n=256,
        mac_k=16,
        workgroup_size_x=128,
        workgroup_size_y=2,
        prefetchInFlight=2,
        prefetchLDSFactor=2,
        types=TypeParameters(
            hgemm["types"],
            trans_A="N",
            trans_B="T",
        ),
    )
    yield mkGEMM(
        hgemm,
        mac_m=128,
        mac_n=256,
        mac_k=16,
        beta=0.0,
        workgroup_size_x=64,
        workgroup_size_y=4,
    )


def codegen():
    yield CodeGenRun(instCount=40000, instructions="comments")
    yield CodeGenRun(instCount=40000, instructions="simple_mi")
    yield CodeGenRun(instCount=40000, instructions="complex_mi_with_coop")


def f16gemm_16x16x32():
    params = dict(
        M=128,
        N=128,
        K=256,
        mac_m=64,
        mac_n=64,
        mac_k=32,
        wave_m=16,
        wave_n=16,
        wave_k=32,
        workgroup_size_x=256,
        workgroup_size_y=1,
    )
    for a, b in product("NT", repeat=2):
        yield GEMMRun(
            **params,
            types=TypeParameters(fp16, trans_A=a, trans_B=b),
        )


def f16gemm_32x32x16():
    params = dict(
        M=256,
        N=256,
        K=128,
        mac_m=128,
        mac_n=128,
        mac_k=16,
        wave_m=32,
        wave_n=32,
        wave_k=16,
        workgroup_size_x=256,
        workgroup_size_y=1,
    )
    for a, b in product("NT", repeat=2):
        yield GEMMRun(
            **params,
            types=TypeParameters(fp16, trans_A=a, trans_B=b),
        )


def f16gemm():
    yield from f16gemm_16x16x32()
    yield from f16gemm_32x32x16()


def f8gemm():
    yield GEMMRun(
        M=1024,
        N=1024,
        K=1024,
        mac_m=64,
        mac_n=16,
        mac_k=64,
        workgroup_size_x=256,
        workgroup_size_y=1,
        types=fp8fp8_fp32,
    )


def mx_gemm_16x16x128_f8f6f4():
    params = dict(
        M=4096,
        N=4096,
        K=8192,
        mac_m=64,
        mac_n=64,
        mac_k=128,
        wave_m=16,
        wave_n=16,
        wave_k=128,
        workgroup_size_x=256,
        workgroup_size_y=1,
    )
    for typeParam in [fp8fp8_fp32, bf8bf8_fp32, fp6fp6_fp32, bf6bf6_fp32, fp4fp4_fp32]:
        for a, b in product("NT", repeat=2):
            yield GEMMRun(
                **params,
                types=TypeParameters(typeParam, trans_A=a, trans_B=b),
            )


def mx_gemm_32x32x64_f8f6f4():
    params = dict(
        M=4096,
        N=4096,
        K=8192,
        mac_m=128,
        mac_n=128,
        mac_k=64,
        wave_m=32,
        wave_n=32,
        wave_k=64,
        workgroup_size_x=256,
        workgroup_size_y=1,
    )
    for typeParam in [fp8fp8_fp32, bf8bf8_fp32, fp6fp6_fp32, bf6bf6_fp32, fp4fp4_fp32]:
        for a, b in product("NT", repeat=2):
            yield GEMMRun(
                **params,
                types=TypeParameters(typeParam, trans_A=a, trans_B=b),
            )


def mx_gemms_f8f6f4():
    yield from mx_gemm_32x32x64_f8f6f4()
    yield from mx_gemm_16x16x128_f8f6f4()


def gemm_mixed_16x16x128_f8f6f4():
    params = dict(
        M=256,
        N=256,
        K=512,
        mac_m=64,
        mac_n=64,
        mac_k=128,
        wave_m=16,
        wave_n=16,
        wave_k=128,
        workgroup_size_x=256,
        workgroup_size_y=1,
    )
    TA = {"fp8", "bf8", "fp6", "bf6", "fp4"}
    TB = {"fp8", "bf8", "fp6", "bf6", "fp4"}
    for A in TA:
        for B in TB:
            yield GEMMRun(
                **params,
                types=TypeParameters(
                    type_A=A,
                    type_B=B,
                    type_C="float",
                    type_D="float",
                    trans_A="T",
                    trans_B="N",
                ),
            )


def gemm_mixed_32x32x64_f8f6f4():
    params = dict(
        M=256,
        N=256,
        K=512,
        mac_m=128,
        mac_n=128,
        mac_k=64,
        wave_m=32,
        wave_n=32,
        wave_k=64,
        workgroup_size_x=256,
        workgroup_size_y=1,
    )
    TA = {"fp8", "bf8", "fp6", "bf6", "fp4"}
    TB = {"fp8", "bf8", "fp6", "bf6", "fp4"}
    for A in TA:
        for B in TB:
            yield GEMMRun(
                **params,
                types=TypeParameters(
                    type_A=A,
                    type_B=B,
                    type_C="float",
                    type_D="float",
                    trans_A="T",
                    trans_B="N",
                ),
            )


def mixedgemm_f8f6f4():
    yield from gemm_mixed_16x16x128_f8f6f4()
    yield from gemm_mixed_32x32x64_f8f6f4()


def _f8f6f4_gemm_macrotiles(
    wave_m, wave_n, wave_k, gemmTypes, mac_m_factor, mac_n_factor, mac_k_factor
):
    params = dict(
        M=256,
        N=256,
        K=512,
        mac_m=wave_m * mac_m_factor,
        mac_n=wave_m * mac_n_factor,
        mac_k=wave_k * mac_k_factor,
        wave_m=wave_m,
        wave_n=wave_n,
        wave_k=wave_k,
        workgroup_size_x=64,
        workgroup_size_y=1,
    )
    yield GEMMRun(**params, types=TypeParameters(gemmTypes, trans_A="T", trans_B="N"))


def gemm_f8f6f4_different_macrotiles():
    for type in [fp8fp8_fp32, fp6fp6_fp32, fp4fp4_fp32, bf8bf8_fp32, bf6bf6_fp32]:
        for m_factor in [2, 4, 8]:
            for n_factor in [2, 4, 8]:
                for k_factor in [2, 4]:
                    yield from _f8f6f4_gemm_macrotiles(
                        16, 16, 128, type, m_factor, n_factor, k_factor
                    )
                    yield from _f8f6f4_gemm_macrotiles(
                        32, 32, 64, type, m_factor, n_factor, k_factor
                    )


def _f8f6f4_gemm_prefetch(wave_m, wave_n, wave_k, gemmTypes, prefetchFactor):
    params = dict(
        M=256,
        N=256,
        K=512,
        mac_m=wave_m * 4,
        mac_n=wave_m * 4,
        mac_k=wave_k * 2,
        wave_m=wave_m,
        wave_n=wave_n,
        wave_k=wave_k,
        workgroup_size_x=256,
        workgroup_size_y=1,
        prefetchInFlight=prefetchFactor,
        prefetchLDSFactor=prefetchFactor,
    )
    yield GEMMRun(
        **params,
        types=TypeParameters(gemmTypes, trans_A="T", trans_B="N"),
    )


def gemm_f8f6f4_test_prefetch():
    # Any factor > 2 fails to generate due to lack of registers/too large imm offset.
    # TODO: test with more aggressive prefetching once we can.
    for factor in [2]:
        for type in [fp8fp8_fp32, fp6fp6_fp32, fp4fp4_fp32, bf8bf8_fp32, bf6bf6_fp32]:
            yield from _f8f6f4_gemm_prefetch(16, 16, 128, type, factor)
            yield from _f8f6f4_gemm_prefetch(32, 32, 64, type, factor)


def bf16gemm_16x16x8():
    yield GEMMRun(
        M=1024,
        N=1024,
        K=1024,
        mac_m=64,
        mac_n=16,
        mac_k=64,
        wave_m=16,
        wave_n=16,
        wave_k=8,
        workgroup_size_x=128,
        workgroup_size_y=1,
        types=bf16_fp32,
    )


def bf16gemm_32x32x4():
    yield GEMMRun(
        M=1024,
        N=1024,
        K=1024,
        mac_m=64,
        mac_n=64,
        mac_k=64,
        wave_m=32,
        wave_n=32,
        wave_k=4,
        workgroup_size_x=128,
        workgroup_size_y=1,
        types=bf16_fp32,
    )


def bf16bf16gemm_16x16x8():
    yield GEMMRun(
        M=1024,
        N=1024,
        K=1024,
        mac_m=64,
        mac_n=16,
        mac_k=64,
        wave_m=16,
        wave_n=16,
        wave_k=8,
        workgroup_size_x=128,
        workgroup_size_y=1,
        types=bf16_bf16,
    )


def bf16bf16gemm_32x32x4():
    yield GEMMRun(
        M=1024,
        N=1024,
        K=1024,
        mac_m=64,
        mac_n=64,
        mac_k=64,
        wave_m=32,
        wave_n=32,
        wave_k=4,
        workgroup_size_x=128,
        workgroup_size_y=1,
        types=bf16_bf16,
    )


def fp4_small():
    yield GEMMRun(
        M=256,
        N=256,
        K=256,
        beta=0.0,
        mac_m=128,
        mac_n=128,
        mac_k=256,
        wave_m=16,
        wave_n=16,
        wave_k=128,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        workgroupMappingDim=0,
        workgroupMappingValue=2,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToLDS",
        loadScale_B="BufferToLDS",
        store="VGPRToGlobalMemoryWithBuffer",
        prefetch=True,
        prefetchInFlight=2,
        prefetchLDSFactor=1,
        prefetchScale=True,
        pretileScale=True,
        swizzleScale=True,
        prefetchMixMemOps=True,
        betaInFma=True,
        scheduler="Priority",
        schedulerCost="LinearWeightedSimple",
        tailLoops=True,
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp4",
            type_B="fp4",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
            scaleSkipPermlane="PreSwizzleScaleGFX950",
        ),
        swizzleTileSize=MKNLTuple(32, 8, 32, 8),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def fp4_small_k():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=256,
        beta=0.0,
        mac_m=128,
        mac_n=128,
        mac_k=256,
        wave_m=16,
        wave_n=16,
        wave_k=128,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        workgroupMappingDim=0,
        workgroupMappingValue=2,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToLDS",
        loadScale_B="BufferToLDS",
        store="VGPRToGlobalMemoryWithBuffer",
        prefetch=True,
        prefetchInFlight=2,
        prefetchLDSFactor=1,
        prefetchScale=True,
        pretileScale=True,
        swizzleScale=True,
        prefetchMixMemOps=True,
        betaInFma=True,
        scheduler="Priority",
        schedulerCost="LinearWeightedSimple",
        tailLoops=True,
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp4",
            type_B="fp4",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
            scaleSkipPermlane="PreSwizzleScaleGFX950",
        ),
        swizzleTileSize=MKNLTuple(32, 8, 32, 8),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def fp4_target():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        mac_m=256,
        mac_n=256,
        mac_k=128,
        wave_m=32,
        wave_n=32,
        wave_k=64,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDSViaVGPR",
        load_B="BufferToLDSViaVGPR",
        loadScale_A="BufferToLDSViaVGPR",
        loadScale_B="BufferToLDSViaVGPR",
        store="VGPRToGlobalMemoryViaLDSWithBuffer",
        prefetch=True,
        prefetchInFlight=2,
        prefetchLDSFactor=2,
        betaInFma=True,
        scheduler="Priority",
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp4",
            type_B="fp4",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
        ),
        swizzleTileSize=MKNLTuple(64, 64 // 32 * 2 * 2, 64, 64 // 32 * 2 * 2),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def fp4_target_d2lds_mi32x32x64_pf2x1():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        mac_m=256,
        mac_n=256,
        mac_k=128,
        wave_m=32,
        wave_n=32,
        wave_k=64,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToLDSViaVGPR",
        loadScale_B="BufferToLDSViaVGPR",
        store="VGPRToGlobalMemoryViaLDSWithBuffer",
        prefetch=True,
        prefetchInFlight=2,
        prefetchLDSFactor=1,
        betaInFma=True,
        scheduler="Priority",
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp4",
            type_B="fp4",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
        ),
        swizzleTileSize=MKNLTuple(64, 64 // 32 * 2 * 2, 64, 64 // 32 * 2 * 2),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def add_wgm(mapping, suite):
    for run in suite:
        run.workgroupMappingDim = mapping[0]
        run.workgroupMappingValue = mapping[1]
        yield run


def addSkipPermlane(suite: Generator[GEMMRun, None, None], value="PreSwizzleScale"):
    for run in suite:
        run.types.scaleSkipPermlane = value
        yield run


def addStreamK(suite: Generator[GEMMRun, None, None], value="Standard"):
    for run in suite:
        run.streamK = value
        if run.types.scale_A == "Separate" or run.types.scale_B == "Separate":
            run.schedulerCost = "LinearWeightedSimpleStreamK"
        yield run


def addLDSBankSwizzle(suite, value="Swizzle"):
    for run in suite:
        run.ldsBankSwizzle = value
        yield run


def fp4_target_d2lds_mi32x32x64_pf2x1_wgm():
    yield from add_wgm((0, 2), fp4_target_d2lds_mi32x32x64_pf2x1())


def fp4_target_d2lds_mi32x32x64_pf4x1_sweep_wgms():
    for wgm_dim in [0, 1]:
        for wgm_value in range(1, 50):
            yield from add_wgm(
                (wgm_dim, wgm_value), fp4_target_d2lds_mi32x32x64_pf2x1()
            )


def fp4_target_d2lds_mi32x32x64_pf4x1():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        beta=0.0,
        mac_m=256,
        mac_n=256,
        mac_k=128,
        wave_m=32,
        wave_n=32,
        wave_k=64,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToVGPR",
        loadScale_B="BufferToVGPR",
        store="VGPRToGlobalMemoryWithBuffer",
        prefetch=True,
        prefetchInFlight=4,
        prefetchLDSFactor=1,
        prefetchScale=True,
        swizzleScale=True,
        prefetchMixMemOps=True,
        betaInFma=True,
        scheduler="Priority",
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp4",
            type_B="fp4",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
        ),
        swizzleTileSize=MKNLTuple(64, 64 // 32 * 2 * 4, 64, 64 // 32 * 2 * 4),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def fp4_target_d2lds_mi32x32x64_st32x8_pf4x1():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        beta=0.0,
        mac_m=128,
        mac_n=128,
        mac_k=256,
        wave_m=32,
        wave_n=32,
        wave_k=64,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToVGPR",
        loadScale_B="BufferToVGPR",
        store="VGPRToGlobalMemoryWithBuffer",
        prefetch=True,
        prefetchInFlight=4,
        prefetchLDSFactor=1,
        prefetchScale=True,
        swizzleScale=True,
        prefetchMixMemOps=True,
        betaInFma=True,
        scheduler="Priority",
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp4",
            type_B="fp4",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
        ),
        swizzleTileSize=MKNLTuple(32, 8, 32, 8),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def fp4_target_d2lds_mi32x32x64_pf4x1_wgm():
    yield from add_wgm((0, 2), fp4_target_d2lds_mi32x32x64_pf4x1())
    yield from add_wgm((0, 2), fp4_target_d2lds_mi32x32x64_st32x8_pf4x1())


def fp4_target_d2lds_mi32x32x64_pf4x1_both():
    yield from fp4_target_d2lds_mi32x32x64_pf4x1()
    yield from fp4_target_d2lds_mi32x32x64_st32x8_pf4x1()
    yield from fp4_target_d2lds_mi32x32x64_pf4x1_wgm()


def fp4_target_d2lds_mi16x16x128_pf4x1():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        beta=0.0,
        mac_m=256,
        mac_n=256,
        mac_k=128,
        wave_m=16,
        wave_n=16,
        wave_k=128,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToVGPR",
        loadScale_B="BufferToVGPR",
        store="VGPRToGlobalMemoryWithBuffer",
        prefetch=True,
        prefetchInFlight=4,
        prefetchLDSFactor=1,
        prefetchScale=True,
        swizzleScale=True,
        prefetchMixMemOps=True,
        betaInFma=True,
        scheduler="Priority",
        schedulerCost="LinearWeightedSimple",
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp4",
            type_B="fp4",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
        ),
        swizzleTileSize=MKNLTuple(64, 128 // 32 * 1 * 4, 64, 128 // 32 * 1 * 4),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def fp4_target_d2lds_mi16x16x128_st32x8_pf4x1():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        beta=0.0,
        mac_m=128,
        mac_n=128,
        mac_k=256,
        wave_m=16,
        wave_n=16,
        wave_k=128,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToVGPR",
        loadScale_B="BufferToVGPR",
        store="VGPRToGlobalMemoryWithBuffer",
        prefetch=True,
        prefetchInFlight=4,
        prefetchLDSFactor=1,
        prefetchScale=True,
        swizzleScale=True,
        prefetchMixMemOps=True,
        betaInFma=True,
        scheduler="Priority",
        schedulerCost="LinearWeightedSimple",
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp4",
            type_B="fp4",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
        ),
        swizzleTileSize=MKNLTuple(32, 8, 32, 8),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def fp4_target_d2lds_mi16x16x128_pf4x1_wgm():
    yield from add_wgm((0, 2), fp4_target_d2lds_mi16x16x128_pf4x1())
    yield from add_wgm((0, 2), fp4_target_d2lds_mi16x16x128_st32x8_pf4x1())


def fp4_target_d2lds_mi16x16x128_pf4x1_both():
    yield from fp4_target_d2lds_mi16x16x128_pf4x1()
    yield from fp4_target_d2lds_mi16x16x128_st32x8_pf4x1()
    yield from fp4_target_d2lds_mi16x16x128_pf4x1_wgm()


def fp4_target_d2lds_mi16x16x128_st32x8_pf2x1():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        beta=0.0,
        mac_m=256,
        mac_n=256,
        mac_k=256,
        wave_m=16,
        wave_n=16,
        wave_k=128,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToLDS",
        loadScale_B="BufferToLDS",
        store="VGPRToGlobalMemoryWithBuffer",
        prefetch=True,
        prefetchInFlight=2,
        prefetchLDSFactor=1,
        prefetchScale=True,
        pretileScale=True,
        swizzleScale=True,
        prefetchMixMemOps=True,
        betaInFma=True,
        scheduler="Priority",
        schedulerCost="LinearWeightedSimple",
        tailLoops=False,
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp4",
            type_B="fp4",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
            scaleSkipPermlane="PreSwizzleScaleGFX950",
        ),
        swizzleTileSize=MKNLTuple(32, 8, 32, 8),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def fp4_mt32x32x256_d2lds_mi16x16x128_st32x8_pf2x1():
    yield GEMMRun(
        M=1024,
        N=1024,
        K=4096,
        beta=0.0,
        mac_m=32,
        mac_n=32,
        mac_k=256,
        wave_m=16,
        wave_n=16,
        wave_k=128,
        wave_b=1,
        workgroup_size_x=64,
        workgroup_size_y=1,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToLDS",
        loadScale_B="BufferToLDS",
        store="VGPRToGlobalMemoryWithBuffer",
        prefetch=True,
        prefetchInFlight=2,
        prefetchLDSFactor=1,
        prefetchScale=True,
        pretileScale=True,
        swizzleScale=True,
        prefetchMixMemOps=True,
        betaInFma=True,
        scheduler="Priority",
        schedulerCost="LinearWeightedSimple",
        matchMemoryAccess=True,
        tailLoops=True,
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp4",
            type_B="fp4",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
            scaleSkipPermlane="PreSwizzleScaleGFX950",
        ),
        swizzleTileSize=MKNLTuple(32, 8, 32, 8),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def fp4_scaled2lds_mi16x16x128_st32x8_pf2x1_sweep():
    """Sweep mac_m, mac_n, and prefetchLDSFactor for fp4_mt32x32x256 configuration."""
    mac_m_values = [32, 64, 96]
    mac_n_values = [32, 64, 96]
    prefetchLDSFactor_values = [1]

    for mac_m in mac_m_values:
        for mac_n in mac_n_values:
            for prefetchLDSFactor in prefetchLDSFactor_values:
                yield GEMMRun(
                    M=1024,
                    N=1024,
                    K=4096,
                    beta=0.0,
                    mac_m=mac_m,
                    mac_n=mac_n,
                    mac_k=256,
                    wave_m=16,
                    wave_n=16,
                    wave_k=128,
                    wave_b=1,
                    workgroup_size_x=64,
                    workgroup_size_y=1,
                    load_A="BufferToLDS",
                    load_B="BufferToLDS",
                    loadScale_A="BufferToLDS",
                    loadScale_B="BufferToLDS",
                    prefetch=True,
                    prefetchInFlight=2,
                    prefetchLDSFactor=prefetchLDSFactor,
                    prefetchScale=True,
                    pretileScale=True,
                    swizzleScale=True,
                    prefetchMixMemOps=True,
                    betaInFma=True,
                    scheduler="Priority",
                    schedulerCost="LinearWeightedSimple",
                    tailLoops=True,
                    types=TypeParameters(
                        trans_A="T",
                        trans_B="N",
                        type_A="fp4",
                        type_B="fp4",
                        type_C="half",
                        type_D="half",
                        type_acc="float",
                        scale_A="Separate",
                        scaleType_A="E8M0",
                        scale_B="Separate",
                        scaleType_B="E8M0",
                        scaleBlockSize=32,
                        scaleSkipPermlane="PreSwizzleScaleGFX950",
                    ),
                    swizzleTileSize=MKNLTuple(32, 8, 32, 8),
                    numOuter=1,
                    numWarmUp=1000,
                    numInner=1000,
                )


def fp4_target_d2lds_mi16x16x128_st32x8_pf2x1_pretile():
    for gemm in fp4_target_d2lds_mi16x16x128_st32x8_pf2x1():
        gemm.prefetchScale = False
        gemm.pretileScale = True
        gemm.loadScale_A = "BufferToLDS"
        gemm.loadScale_B = "BufferToLDS"
        gemm.swizzleTileSize = MKNLTuple(64, 4, 64, 4)
        gemm.types.scaleSkipPermlane = "PreSwizzleScale"
        yield gemm


def fp4_target_d2lds_mi16x16x128_st32x8_pf2x1_pad():
    for gemm in fp4_target_d2lds_mi16x16x128_st32x8_pf2x1():
        gemm.loadScale_A = "BufferToLDS"
        gemm.loadScale_B = "BufferToLDS"
        gemm.swizzleTileSize = MKNLTuple(64, 4, 64, 4)
        gemm.types.scaleSkipPermlane = "PreSwizzleScale"
        gemm.padLDS_A = (2048, 8)
        gemm.padLDS_B = (2048, 8)
        yield gemm


def fp4_target_mxd2lds_mi16x16x128_st32x8_pf2x1():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        beta=0.0,
        mac_m=128,
        mac_n=128,
        mac_k=256,
        wave_m=16,
        wave_n=16,
        wave_k=128,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToLDS",
        loadScale_B="BufferToLDS",
        store="VGPRToGlobalMemoryWithBuffer",
        prefetch=True,
        prefetchInFlight=2,
        prefetchLDSFactor=1,
        prefetchScale=True,
        swizzleScale=True,
        prefetchMixMemOps=True,
        betaInFma=True,
        scheduler="Priority",
        schedulerCost="LinearWeightedSimple",
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp4",
            type_B="fp4",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
            scaleSkipPermlane="PreSwizzleScale",
        ),
        swizzleTileSize=MKNLTuple(32, 8, 32, 8),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def fp4_target_d2lds_mi16x16x128_st32x8_pf2x1_wgm():
    yield from add_wgm((0, 2), fp4_target_d2lds_mi16x16x128_st32x8_pf2x1())


def fp4_target_d2lds_mi16x16x128_st32x8_pf2x1_both():
    yield from fp4_target_d2lds_mi16x16x128_st32x8_pf2x1()
    yield from fp4_target_d2lds_mi16x16x128_st32x8_pf2x1_wgm()


def fp4_target_d2lds_mi16x16x128_st32x8_pf2x1_wgm_lds_swizzle():
    yield from addLDSBankSwizzle(fp4_target_d2lds_mi16x16x128_st32x8_pf2x1_wgm())


def does_this_fail():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        beta=0.0,
        mac_m=64,
        mac_n=64,
        mac_k=128,
        wave_m=16,
        wave_n=16,
        wave_k=128,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToVGPR",
        loadScale_B="BufferToVGPR",
        store="VGPRToGlobalMemoryWithBuffer",
        prefetch=True,
        prefetchInFlight=4,
        prefetchLDSFactor=1,
        prefetchScale=False,
        swizzleScale=False,
        prefetchMixMemOps=True,
        betaInFma=True,
        scheduler="Priority",
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp4",
            type_B="fp4",
            type_C="half",
            type_D="half",
            type_acc="float",
        ),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def fp4_single_scale_target_d2lds_mi16x16x128_pf4x1():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        beta=0.0,
        mac_m=256,
        mac_n=256,
        mac_k=128,
        wave_m=16,
        wave_n=16,
        wave_k=128,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToVGPR",
        loadScale_B="BufferToVGPR",
        store="VGPRToGlobalMemoryWithBuffer",
        prefetch=True,
        prefetchInFlight=4,
        prefetchLDSFactor=1,
        prefetchScale=False,
        swizzleScale=False,
        prefetchMixMemOps=True,
        betaInFma=True,
        scaleValue_A=1e-2,
        scaleValue_B=1e-2,
        scheduler="Priority",
        schedulerCost="LinearWeightedSimple",
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp4",
            type_B="fp4",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="SingleScale",
            scaleType_A="E8M0",
            scale_B="SingleScale",
            scaleType_B="E8M0",
        ),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def fp4_single_scale_target_d2lds_mi16x16x128_pf4x1_wgm():
    yield from add_wgm((0, 2), fp4_single_scale_target_d2lds_mi16x16x128_pf4x1())


def fp4_d2lds_wgts256x256x256():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        beta=0.0,
        mac_m=256,
        mac_n=256,
        mac_k=256,
        wave_m=16,
        wave_n=16,
        wave_k=128,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToLDS",
        loadScale_B="BufferToLDS",
        store="VGPRToGlobalMemoryWithBuffer",
        prefetch=True,
        prefetchInFlight=2,
        prefetchLDSFactor=1,
        prefetchScale=True,
        swizzleScale=True,
        prefetchMixMemOps=True,
        scheduler="Priority",
        schedulerCost="LinearWeightedSimple",
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp4",
            type_B="fp4",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleSkipPermlane="PreSwizzleScale",
        ),
        swizzleTileSize=MKNLTuple(64, 8, 64, 8),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def fp4_kernels_no_wgm():
    yield from fp4_target()
    yield from fp4_target_d2lds_mi32x32x64_pf2x1()
    yield from fp4_target_d2lds_mi32x32x64_pf4x1()
    yield from fp4_target_d2lds_mi32x32x64_st32x8_pf4x1()
    yield from fp4_target_d2lds_mi16x16x128_pf4x1()
    yield from fp4_target_d2lds_mi16x16x128_st32x8_pf4x1()
    yield from fp4_small()
    yield from fp4_small_k()
    # yield from fp4_single_scale_target_d2lds_mi16x16x128_pf4x1()


def fp4_kernels_wgm():
    yield from fp4_target_d2lds_mi32x32x64_pf2x1_wgm()
    yield from fp4_target_d2lds_mi32x32x64_pf4x1_wgm()
    yield from fp4_target_d2lds_mi16x16x128_pf4x1_wgm()
    yield from fp4_single_scale_target_d2lds_mi16x16x128_pf4x1_wgm()


def fp4_kernels_wgm_streamk():
    yield from addStreamK(fp4_kernels_wgm())


def fp4_16x16x128_scale_options():
    yield from fp4_target_d2lds_mi16x16x128_pf4x1_wgm()
    yield from addSkipPermlane(fp4_target_d2lds_mi16x16x128_pf4x1_wgm())
    yield from fp4_target_d2lds_mi16x16x128_st32x8_pf2x1_wgm()
    yield from addSkipPermlane(fp4_target_d2lds_mi16x16x128_st32x8_pf2x1_wgm())


def fp4_32x32x64_scale_options():
    yield from fp4_target_d2lds_mi32x32x64_pf4x1_wgm()
    yield from addSkipPermlane(fp4_target_d2lds_mi32x32x64_pf4x1_wgm())


def fp4_kernels():
    yield from fp4_kernels_no_wgm()
    yield from fp4_kernels_wgm()
    yield from fp4_kernels_wgm_streamk()
    yield from fp4_16x16x128_scale_options()
    yield from fp4_32x32x64_scale_options()
    yield from fp4_d2lds_wgts256x256x256()
    yield from fp4_target_d2lds_mi16x16x128_st32x8_pf2x1_pretile()
    yield from fp4_target_d2lds_mi16x16x128_st32x8_pf2x1_pad()
    yield from test_streamk_fp4()


def fp4_kernels_streamk():
    """Only the FP4 kernel configurations that pass with StreamK enabled."""
    yield from addStreamK(fp4_target())
    yield from addStreamK(fp4_target_d2lds_mi32x32x64_pf2x1())
    yield from addStreamK(fp4_target_d2lds_mi32x32x64_pf4x1())
    yield from addStreamK(fp4_target_d2lds_mi16x16x128_pf4x1())
    yield from addStreamK(fp4_d2lds_wgts256x256x256())


def test_streamk_fp4():
    kernel = fp4_target_d2lds_mi16x16x128_st32x8_pf2x1()
    kernel = addStreamK(kernel)
    for k in kernel:
        k.M = 128
        k.N = 128
        k.K = 76800
        k.mac_m = (64,)
        k.mac_n = (32,)
        k.mac_k = (256,)
        k.prefetchInFlight = 2

        yield k


def fp4_target_sweep_wgms():
    for wgm_dim in [0, 1]:
        for wgm_value in range(1, 50):
            yield from add_wgm(
                (wgm_dim, wgm_value), fp4_single_scale_target_d2lds_mi16x16x128_pf4x1()
            )


def generate_gfx950():
    yield from fp4_kernels()


def mxfp8_target_128x256():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        mac_m=128,
        mac_n=256,
        mac_k=128,
        wave_m=16,
        wave_n=16,
        wave_k=128,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDSViaVGPR",
        load_B="BufferToLDSViaVGPR",
        loadScale_A="BufferToLDSViaVGPR",
        loadScale_B="BufferToLDSViaVGPR",
        store="VGPRToGlobalMemoryViaLDSWithBuffer",
        prefetch=True,
        prefetchInFlight=2,
        prefetchLDSFactor=2,
        betaInFma=True,
        scheduler="Priority",
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp8",
            type_B="fp8",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
        ),
        swizzleTileSize=MKNLTuple(64, 128 // 32 * 1 * 2, 64, 128 // 32 * 1 * 2),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def mxfp8_target_256x128():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        mac_m=256,
        mac_n=128,
        mac_k=128,
        wave_m=16,
        wave_n=16,
        wave_k=128,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDSViaVGPR",
        load_B="BufferToLDSViaVGPR",
        loadScale_A="BufferToLDSViaVGPR",
        loadScale_B="BufferToLDSViaVGPR",
        store="VGPRToGlobalMemoryViaLDSWithBuffer",
        prefetch=True,
        prefetchInFlight=2,
        prefetchLDSFactor=2,
        betaInFma=True,
        scheduler="Priority",
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp8",
            type_B="fp8",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
        ),
        swizzleTileSize=MKNLTuple(64, 128 // 32 * 1 * 2, 64, 128 // 32 * 1 * 2),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def mxfp8_target_d2lds_mi32x32x64_pf2x1():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        mac_m=256,
        mac_n=256,
        mac_k=128,
        wave_m=32,
        wave_n=32,
        wave_k=64,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToLDSViaVGPR",
        loadScale_B="BufferToLDSViaVGPR",
        store="VGPRToGlobalMemoryViaLDSWithBuffer",
        prefetch=True,
        prefetchInFlight=2,
        prefetchLDSFactor=1,
        betaInFma=True,
        scheduler="Priority",
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp8",
            type_B="fp8",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
        ),
        swizzleTileSize=MKNLTuple(64, 64 // 32 * 2 * 2, 64, 64 // 32 * 2 * 2),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def mxfp8_target_d2lds_mi32x32x64_pf2x1_wgm():
    yield from add_wgm((0, 4), mxfp8_target_d2lds_mi32x32x64_pf2x1())


def mxfp8_target_d2lds_mi32x32x64_pf4x1():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        beta=0.0,
        mac_m=128,
        mac_n=128,
        mac_k=128,
        wave_m=32,
        wave_n=32,
        wave_k=64,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToVGPR",
        loadScale_B="BufferToVGPR",
        store="VGPRToGlobalMemoryWithBuffer",
        prefetch=True,
        prefetchInFlight=4,
        prefetchLDSFactor=1,
        prefetchScale=True,
        swizzleScale=True,
        prefetchMixMemOps=True,
        betaInFma=True,
        scheduler="Priority",
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp8",
            type_B="fp8",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
        ),
        swizzleTileSize=MKNLTuple(64, 64 // 32 * 2 * 4, 64, 64 // 32 * 2 * 4),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def mxfp8_target_d2lds_mi32x32x64_pf4x1_wgm():
    yield from add_wgm((1, 2), mxfp8_target_d2lds_mi32x32x64_pf4x1())


def mxfp8_target_d2lds_mi16x16x128_pf4x1():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        beta=0.0,
        mac_m=128,
        mac_n=128,
        mac_k=128,
        wave_m=16,
        wave_n=16,
        wave_k=128,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToVGPR",
        loadScale_B="BufferToVGPR",
        store="VGPRToGlobalMemoryWithBuffer",
        prefetch=True,
        prefetchInFlight=4,
        prefetchLDSFactor=1,
        prefetchScale=True,
        swizzleScale=True,
        prefetchMixMemOps=True,
        betaInFma=True,
        scheduler="Priority",
        schedulerCost="LinearWeightedSimple",
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp8",
            type_B="fp8",
            type_C="half",
            type_D="half",
            type_acc="float",
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
        ),
        swizzleTileSize=MKNLTuple(64, 128 // 32 * 1 * 4, 64, 128 // 32 * 1 * 4),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def mxfp8_target_d2lds_mi16x16x128_pf4x1_wgm():
    yield from add_wgm((1, 1), mxfp8_target_d2lds_mi16x16x128_pf4x1())


def fp8_target_d2lds_mi16x16x128_pf4x1():
    yield GEMMRun(
        M=4096,
        N=4096,
        K=32768,
        beta=0.0,
        mac_m=128,
        mac_n=128,
        mac_k=128,
        wave_m=16,
        wave_n=16,
        wave_k=128,
        wave_b=1,
        workgroup_size_x=128,
        workgroup_size_y=2,
        load_A="BufferToLDS",
        load_B="BufferToLDS",
        loadScale_A="BufferToVGPR",
        loadScale_B="BufferToVGPR",
        store="VGPRToGlobalMemoryWithBuffer",
        prefetch=True,
        prefetchInFlight=4,
        prefetchLDSFactor=1,
        prefetchScale=False,
        swizzleScale=False,
        prefetchMixMemOps=True,
        betaInFma=True,
        scheduler="Priority",
        types=TypeParameters(
            trans_A="T",
            trans_B="N",
            type_A="fp8",
            type_B="fp8",
            type_C="half",
            type_D="half",
            type_acc="float",
        ),
        numOuter=1,
        numWarmUp=1000,
        numInner=1000,
    )


def fp8_target_d2lds_mi16x16x128_pf4x1_wgm():
    yield from add_wgm((1, 2), fp8_target_d2lds_mi16x16x128_pf4x1())


def mxfp8_kernels_no_wgm():
    yield from mxfp8_target_128x256()
    yield from mxfp8_target_256x128()
    yield from mxfp8_target_d2lds_mi32x32x64_pf2x1()
    yield from mxfp8_target_d2lds_mi32x32x64_pf4x1()
    yield from mxfp8_target_d2lds_mi16x16x128_pf4x1()


def mxfp8_kernels_wgm():
    yield from mxfp8_target_d2lds_mi32x32x64_pf2x1_wgm()
    yield from mxfp8_target_d2lds_mi32x32x64_pf4x1_wgm()
    yield from mxfp8_target_d2lds_mi16x16x128_pf4x1_wgm()


def mxfp8_16x16x128_scale_options():
    yield from addSkipPermlane(mxfp8_target_d2lds_mi16x16x128_pf4x1_wgm())


def mxfp8_32x32x64_scale_options():
    yield from addSkipPermlane(mxfp8_target_d2lds_mi32x32x64_pf4x1_wgm())


def mxfp8_target_sweep_wgms():
    for wgm_dim in [0, 1]:
        for wgm_value in range(
            1, 9
        ):  # All above *fp8*wgm optimized only for this sweep range
            yield from add_wgm(
                (wgm_dim, wgm_value), fp8_target_d2lds_mi16x16x128_pf4x1()
            )


def mxfp8_kernels():
    yield from mxfp8_kernels_no_wgm()
    yield from mxfp8_kernels_wgm()
    yield from mxfp8_16x16x128_scale_options()
    yield from mxfp8_32x32x64_scale_options()


def fp8_kernels():
    yield from fp8_target_d2lds_mi16x16x128_pf4x1()
    yield from fp8_target_d2lds_mi16x16x128_pf4x1_wgm()


def all():
    if rocm_gfx().startswith("gfx95"):
        yield from fp4_kernels()
        yield from fp8_kernels()
        yield from mxfp8_kernels()
        yield from mx_gemms_f8f6f4()

    yield from sgemm()
    yield from hgemm()
    yield from hgemm_no_store_LDS()
    yield from streamk()
    yield from streamk_sweep()
    yield from scalar_is_zero()
    yield from smallMN_largeK_fp32()
    yield from codegen()


def all_gfx120X():
    yield from hgemm_gfx120X()
    yield from codegen()


def hgemm_guideposts():
    yield from guidepost_1()
    yield from guidepost_2()


def priority_problems():
    return {
        "1. HGEMM Guidepost": {
            "M": 7680,
            "N": 8448,
            "trans_A": "N",
            "trans_B": "T",
        },
        "2. Halfs": {"type_A": "half"},
        "3. Floats": {"type_A": "float"},
    }