TheRock/rocm-libraries/projects/composablekernel/example/ck_tile/49_sageattention/generate.py

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

# generate kernel instances to speed up compilation

import argparse
import sys
import importlib.util
from enum import IntEnum
from pathlib import Path
import pkgutil
from typing import List, Optional

import codegen.ops
from codegen.cmake_config import GEN_DIR


class HandlerId(IntEnum):
    LIST_BLOBS = 0
    WRITE_BLOBS = 1


def _load_ops_module(importer, module_name: str):
    """Load a submodule from codegen.ops, matching Loader.load_module behavior."""
    spec = importer.find_spec(module_name)
    loader = spec.loader
    if hasattr(loader, "load_module"):
        return loader.load_module(module_name)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    loader.exec_module(module)
    return module


# inspect all modules under 'codegen.ops' and register API handlers
ops = []
for importer, module_name, _ in pkgutil.iter_modules(codegen.ops.__path__):
    full_module_name = "%s.%s" % (codegen.ops.__name__, module_name)
    ops.append(_load_ops_module(importer, module_name))
# Strip "sageattn_" so module sageattn_fwd registers as CLI key "fwd".
unwanted_prefix = "sageattn_"
handlers = dict(
    [
        (
            (
                op.__name__[len(unwanted_prefix) :]
                if op.__name__.startswith(unwanted_prefix)
                else op.__name__
            ),
            (op.list_blobs, op.write_blobs),
        )
        for op in ops
    ]
)
assert 0 < len(handlers)


def write_blobs(
    targets: List[str],
    output_dir: Optional[str],
    api_list: List[str],
    filters_list: List[str],
    optdim_list: List[int],
    receipt,
    mask_impl,
) -> None:
    if output_dir is None:
        output_dir = Path(__file__).parent
    else:
        output_dir = Path(output_dir) / GEN_DIR

    output_dir.mkdir(parents=True, exist_ok=True)

    for api, kernel_filter in zip(api_list, filters_list):
        handler = handlers[api][HandlerId.WRITE_BLOBS]
        handler(targets, output_dir, kernel_filter, receipt, optdim_list, mask_impl)


# list all the files that will be generated
def list_blobs(
    targets: List[str],
    output_file: Optional[str],
    api_list: List[str],
    filters_list: List[str],
    optdim_list: List[int],
    receipt,
    mask_impl,
) -> None:
    assert output_file is not None
    file_path = Path(output_file)

    # create an empty file / drop its contents if it exists
    open(file_path, "w").close()

    for api, kernel_filter in zip(api_list, filters_list):
        handler = handlers[api][HandlerId.LIST_BLOBS]
        handler(targets, file_path, kernel_filter, receipt, optdim_list, mask_impl)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog="generate",
        description="Generate SageAttention CK_tile kernel/API blobs.",
    )
    parser.add_argument(
        "--targets",
        default="gfx9,gfx950",
        required=False,
        help="list of GPU targets, separated by comma.",
    )
    parser.add_argument(
        "-a",
        "--api",
        default="fwd",
        required=False,
        help="Codegen API key(s), comma-separated (e.g. fwd -> module codegen.ops.sageattn_fwd).",
    )
    parser.add_argument(
        "-o",
        "--output_dir",
        required=False,
        help="write all the blobs into a directory",
    )
    parser.add_argument(
        "-l", "--list_blobs", required=False, help="list all the kernels to a file"
    )
    # TODO: if using filter, must apply same value to output_dir and list_blobs
    parser.add_argument(
        "-f",
        "--filter",
        default="",
        required=False,
        help="filter out kernels that need to generate, using fnmatch module",
    )

    parser.add_argument(
        "-m",
        "--mask",
        default="simplified",
        required=False,
        help="mask implementation, simplified/generic",
    )

    parser.add_argument(
        "-r",
        "--receipt",
        default=0,
        required=False,
        help="Codegen receipt index. SageAttention forward currently uses receipt 0 only; "
        "the value is passed through to ops (see get_product in sageattn_fwd.py).",
    )

    parser.add_argument(
        "--optdim",
        default="-1",
        required=False,
        help="only optimize the hdim in the list. separated by comma. -1 is the default choice. "
        "e.g. --optdim=32,64,128,256",
    )

    args = parser.parse_args()
    targets = args.targets.split(",")
    api_list = args.api.split(",")
    filter_list = args.filter.split(",")
    filter_list.extend([""] * (len(api_list) - len(filter_list)))
    optdim_list = [int(hdim) for hdim in args.optdim.split(",")]

    if args.list_blobs is not None:
        list_blobs(
            targets,
            args.list_blobs,
            api_list,
            filter_list,
            optdim_list,
            int(args.receipt),
            mask_impl=args.mask,
        )
    else:
        write_blobs(
            targets,
            args.output_dir,
            api_list,
            filter_list,
            optdim_list,
            int(args.receipt),
            mask_impl=args.mask,
        )