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

################################################################################
#
# MIT License
#
# Copyright 2024-2025 AMD ROCm(TM) Software
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell cop-
# ies of the Software, and to permit persons to whom the Software is furnished
# to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IM-
# PLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNE-
# CTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#
################################################################################

import itertools
from pathlib import Path

from rrperf.problems import TypeParameters
from rrperf.rrsuites import codegen, mkGEMM

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

TYPE_CONFIGS_F8 = {
    "bf8bf8_fp16": dict(
        type_A="bf8",
        type_B="bf8",
        type_C="half",
        type_D="half",
        type_acc="half",
        wave_k=64,
    ),
    # TODO: support these types (unit_gfx1250 suite)
    # "bf8fp8_fp16": dict(
    #     type_A="bf8",
    #     type_B="fp8",
    #     type_C="half",
    #     type_D="half",
    #     type_acc="half",
    #     wave_k=64,
    # ),
    # "fp8bf8_fp16": dict(
    #     type_A="fp8",
    #     type_B="bf8",
    #     type_C="half",
    #     type_D="half",
    #     type_acc="half",
    #     wave_k=64,
    # ),
    "fp8fp8_fp16": dict(
        type_A="fp8",
        type_B="fp8",
        type_C="half",
        type_D="half",
        type_acc="half",
        wave_k=64,
    ),
    "bf8bf8_fp32": dict(
        type_A="bf8",
        type_B="bf8",
        type_C="float",
        type_D="float",
        type_acc="float",
        wave_k=64,
    ),
    "bf8fp8_fp32": dict(
        type_A="bf8",
        type_B="fp8",
        type_C="float",
        type_D="float",
        type_acc="float",
        wave_k=64,
    ),
    "fp8bf8_fp32": dict(
        type_A="fp8",
        type_B="bf8",
        type_C="float",
        type_D="float",
        type_acc="float",
        wave_k=64,
    ),
    "fp8fp8_fp32": dict(
        type_A="fp8",
        type_B="fp8",
        type_C="float",
        type_D="float",
        type_acc="float",
        wave_k=64,
    ),
}
TYPE_CONFIGS_F16 = {
    "bf16_bf16": dict(
        type_A="bf16",
        type_B="bf16",
        type_C="bf16",
        type_D="bf16",
        type_acc="bf16",
        wave_k=32,
    ),
    "fp16_fp16": dict(
        type_A="half",
        type_B="half",
        type_C="half",
        type_D="half",
        type_acc="half",
        wave_k=32,
    ),
    "fp16_fp32": dict(
        type_A="half",
        type_B="half",
        type_C="float",
        type_D="float",
        type_acc="float",
        wave_k=32,
    ),
    "bf16_fp32": dict(
        type_A="bf16",
        type_B="bf16",
        type_C="float",
        type_D="float",
        type_acc="float",
        wave_k=32,
    ),
}
TYPE_CONFIGS_F32 = {
    "fp32_fp32": dict(
        type_A="float",
        type_B="float",
        type_C="float",
        type_D="float",
        type_acc="float",
        wave_k=4,
    ),
}

TYPE_CONFIGS_MX = dict()
for a in ["fp8", "bf8", "fp6", "bf6", "fp4"]:
    for b in ["fp8", "bf8", "fp6", "bf6", "fp4"]:
        TYPE_CONFIGS_MX[f"{a}_{b}_fp32"] = dict(
            type_A=a,
            type_B=b,
            type_C="float",
            type_D="float",
            type_acc="float",
            wave_k=128,
        )

BASE_PROBLEM_CONFIGS = {
    "GEMM_512x512x512": dict(
        M=512,
        N=512,
        K=512,
    ),
}

BASE_SOLUTION_CONFIGS = {
    "1250_default": dict(
        mac_m=64,
        mac_n=64,
        mac_k=128,
        wave_m=16,
        wave_n=16,
        workgroup_size_x=64,
        workgroup_size_y=2,
    )
}


def merge_dicts(dicts: list[dict]) -> dict:
    merged = {}
    for d in dicts:
        merged.update(d)
    return merged


# Takes a dict of parameters, extracts those that belong to TypeParameters
# Returns a the input dict with a new types field containing a TypeParameters object
def create_type_parameters(params: dict) -> dict:
    type_fields = {f.name for f in TypeParameters.__dataclass_fields__.values()}
    type_params_dict = {
        k: params.pop(k) for k in list(params.keys()) if k in type_fields
    }
    params["types"] = TypeParameters(**type_params_dict)
    return params


# Run all combinations of the configurations in a grid list, with optional filter
# If parameters overlap, later dicts overwrite earlier ones
def iterate_grid(config_grid: list[list[dict]], skip_filter: callable = None):
    for config_combo in itertools.product(*config_grid):
        params = merge_dicts(config_combo)
        params = create_type_parameters(params)
        if skip_filter is not None and skip_filter(params):
            continue
        yield mkGEMM(params)


def unit_gfx1250():
    base_configs = [
        merge_dicts(
            [
                BASE_PROBLEM_CONFIGS["GEMM_512x512x512"],
                BASE_SOLUTION_CONFIGS["1250_default"],
                dict(numWarmUp=1, numOuter=1, numInner=1),
            ]
        )
    ]
    type_configs_half_float = list(TYPE_CONFIGS_F32.values()) + list(
        TYPE_CONFIGS_F16.values()
    )
    mac_configs_half_float = [
        dict(mac_m=64, mac_n=64, mac_k=64),
    ]
    type_configs_remaining = list(TYPE_CONFIGS_F8.values()) + list(
        TYPE_CONFIGS_MX.values()
    )
    yield from iterate_grid(
        [base_configs, type_configs_half_float, mac_configs_half_float]
    )
    yield from iterate_grid([base_configs, type_configs_remaining])


def sgemm_gfx1250():
    base_configs = [
        merge_dicts(
            [
                BASE_PROBLEM_CONFIGS["GEMM_512x512x512"],
                BASE_SOLUTION_CONFIGS["1250_default"],
                dict(mac_m=64, mac_n=64, mac_k=64),
            ]
        )
    ]
    scalar_configs = [
        dict(beta=0.0),
        dict(beta=0.5),
    ]
    type_configs = list(TYPE_CONFIGS_F32.values())
    mac_configs = [
        dict(mac_m=64, mac_n=64, mac_k=64),
        dict(mac_m=128, mac_n=64, mac_k=16),
    ]
    yield from iterate_grid([base_configs, type_configs, scalar_configs])
    yield from iterate_grid([base_configs, type_configs, mac_configs])


def hgemm_gfx1250():
    base_configs = [
        merge_dicts(
            [
                BASE_PROBLEM_CONFIGS["GEMM_512x512x512"],
                BASE_SOLUTION_CONFIGS["1250_default"],
                dict(mac_m=64, mac_n=64, mac_k=64),
            ]
        )
    ]
    scalar_configs = [
        dict(beta=0.0),
        dict(beta=0.5),
    ]
    type_configs = list(TYPE_CONFIGS_F16.values())
    mac_configs = [
        dict(mac_m=64, mac_n=64, mac_k=64),
        # TODO: Support this mac config
        # dict(mac_m=128, mac_n=256, mac_k=32),
    ]
    trans_configs = [
        dict(trans_A="N", trans_B="N"),
        dict(trans_A="T", trans_B="N"),
        dict(trans_A="T", trans_B="T"),
        dict(trans_A="N", trans_B="T"),
    ]
    workgroup_configs = [
        dict(workgroup_size_x=128, workgroup_size_y=1),
        dict(workgroup_size_x=64, workgroup_size_y=2),
    ]
    prefetch_configs = [
        dict(prefetch=True, prefetchInFlight=2, prefetchLDSFactor=2),
        dict(prefetch=True, prefetchInFlight=2, prefetchLDSFactor=0),
    ]
    scheduler_configs = [
        dict(scheduler="Priority"),
        dict(scheduler="Cooperative"),
        dict(scheduler="Sequential"),
    ]
    yield from iterate_grid([base_configs, type_configs, scalar_configs])
    yield from iterate_grid([base_configs, type_configs, mac_configs])
    yield from iterate_grid([base_configs, type_configs, trans_configs])
    yield from iterate_grid([base_configs, type_configs, workgroup_configs])
    yield from iterate_grid([base_configs, type_configs, prefetch_configs])
    yield from iterate_grid([base_configs, type_configs, scheduler_configs])


def f8gemm_gfx1250():
    base_configs = [
        merge_dicts(
            [
                BASE_PROBLEM_CONFIGS["GEMM_512x512x512"],
                BASE_SOLUTION_CONFIGS["1250_default"],
            ]
        )
    ]
    type_configs = list(TYPE_CONFIGS_F8.values())
    workgroupMapping_configs = [
        dict(workgroupMappingDim=1, workgroupMappingValue=2),
        dict(workgroupMappingDim=-1, workgroupMappingValue=-1),
    ]
    trans_configs = [
        dict(trans_A="N", trans_B="N"),
        dict(trans_A="T", trans_B="N"),
        dict(trans_A="T", trans_B="T"),
        dict(trans_A="N", trans_B="T"),
    ]
    prefetch_configs = [
        dict(
            prefetch=True,
            prefetchInFlight=4,
            prefetchLDSFactor=1,
            prefetchMixMemOps=True,
        ),
        dict(prefetch=False),
    ]
    # direct2lds_configs = [
    #     dict(direct2LDS_A=True, direct2LDS_B=True),
    #     dict(direct2LDS_A=False, direct2LDS_B=False),
    # ]

    yield from iterate_grid([base_configs, type_configs, workgroupMapping_configs])
    yield from iterate_grid([base_configs, type_configs, trans_configs])
    yield from iterate_grid([base_configs, type_configs, prefetch_configs])
    # yield from iterate_grid([base_configs, type_configs, direct2lds_configs])


def f8f6f4_gemm_gfx1250():
    base_configs = [
        merge_dicts(
            [
                BASE_PROBLEM_CONFIGS["GEMM_512x512x512"],
                BASE_SOLUTION_CONFIGS["1250_default"],
            ]
        )
    ]
    type_configs = list(TYPE_CONFIGS_MX.values())
    workgroupMapping_configs = [
        dict(workgroupMappingDim=0, workgroupMappingValue=2),
        dict(workgroupMappingDim=-1, workgroupMappingValue=-1),
    ]
    trans_configs = [
        dict(trans_A="N", trans_B="N"),
        dict(trans_A="T", trans_B="N"),
        dict(trans_A="T", trans_B="T"),
        dict(trans_A="N", trans_B="T"),
    ]
    # prefetch_configs = [
    #     dict(
    #         prefetch=True,
    #         prefetchInFlight=4,
    #         prefetchLDSFactor=1,
    #         prefetchMixMemOps=True,
    #     ),
    #     dict(prefetch=False),
    # ]
    # direct2lds_configs = [
    #     dict(direct2LDS_A=True, direct2LDS_B=True),
    #     dict(direct2LDS_A=False, direct2LDS_B=False),
    # ]
    scale_type_configs = [
        dict(
            scale_A="None",
            scaleType_A="None",
            scale_B="None",
            scaleType_B="None",
            scaleBlockSize=-1,
        ),
        dict(
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
        ),
    ]
    lds_scale_configs = [
        # TODO: Support this config
        # dict(loadLDSScale_A=True, loadLDSScale_B=True),
        dict(loadLDSScale_A=False, loadLDSScale_B=False),
    ]
    scale_swizzle_prefetch_configs = [
        dict(prefetchScale=True, swizzleScale=True, storeLDS_D=False),
        dict(prefetchScale=False, swizzleScale=False, storeLDS_D=True),
    ]
    skip_permlane_configs = [
        # TODO: Support this config
        # dict(scaleSkipPermlane=True),
        dict(scaleSkipPermlane=False),
    ]
    scale_ON_config = [
        dict(
            scale_A="Separate",
            scaleType_A="E8M0",
            scale_B="Separate",
            scaleType_B="E8M0",
            scaleBlockSize=32,
        ),
    ]
    trans_TN_config = [
        dict(trans_A="T", trans_B="N"),
    ]

    yield from iterate_grid([base_configs, type_configs, workgroupMapping_configs])
    yield from iterate_grid([base_configs, type_configs, trans_configs])
    # TODO: support these configs
    # yield from iterate_grid([base_configs, type_configs, prefetch_configs])
    # yield from iterate_grid([base_configs, type_configs, direct2lds_configs])
    yield from iterate_grid([base_configs, type_configs, scale_type_configs])
    yield from iterate_grid(
        [
            base_configs,
            type_configs,
            trans_TN_config,
            scale_ON_config,
            lds_scale_configs,
            scale_swizzle_prefetch_configs,
        ],
    )
    yield from iterate_grid(
        [
            base_configs,
            type_configs,
            trans_TN_config,
            scale_ON_config,
            skip_permlane_configs,
        ]
    )


def stream_k_gfx1250(sweep=False):
    base_configs = [
        merge_dicts(
            [
                BASE_PROBLEM_CONFIGS["GEMM_512x512x512"],
                BASE_SOLUTION_CONFIGS["1250_default"],
                dict(mac_m=64, mac_n=64, mac_k=64),
            ]
        )
    ]
    type_configs = list(TYPE_CONFIGS_F16.values()) + list(TYPE_CONFIGS_F32.values())

    if sweep:
        mac_configs = [
            dict(mac_m=mac_m, mac_n=mac_n, mac_k=mac_k)
            for mac_m in [64, 128]
            for mac_n in [64, 128, 256]
            for mac_k in [32, 64]
        ]
    else:
        mac_configs = [
            dict(mac_m=128, mac_n=256, mac_k=32),
        ]
    stream_k_configs = [
        dict(streamK=True, streamKTwoTile=False),
        dict(streamK=True, streamKTwoTile=True),
    ]
    trans_configs = [
        dict(trans_A="N", trans_B="T"),
    ]

    def skip_filter(params):
        return params["streamKTwoTile"] and (
            params["mac_m"] * params["mac_n"] * params["mac_k"] >= (64 * 256 * 64)
        )

    yield from iterate_grid(
        [base_configs, type_configs, trans_configs, mac_configs, stream_k_configs],
        skip_filter=skip_filter,
    )


def all_gfx1250():
    yield from sgemm_gfx1250()
    yield from hgemm_gfx1250()
    yield from f8gemm_gfx1250()
    yield from f8f6f4_gemm_gfx1250()
    # TODO: support stream k
    # yield from stream_k_gfx1250(sweep=True)
    yield from codegen()