From 3aefb560e02d3bc27f7bb05f24bd3c0c79731032 Mon Sep 17 00:00:00 2001 From: "PoYen, Chen" Date: Mon, 8 Jul 2024 07:17:24 +0000 Subject: [PATCH 1/2] Remove "EXAMPLE_" prefix of cmake variables --- example/ck_tile/01_fmha/CMakeLists.txt | 28 +++++++++++++------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/example/ck_tile/01_fmha/CMakeLists.txt b/example/ck_tile/01_fmha/CMakeLists.txt index 241213a6ec..ad11bd5bc7 100644 --- a/example/ck_tile/01_fmha/CMakeLists.txt +++ b/example/ck_tile/01_fmha/CMakeLists.txt @@ -1,27 +1,27 @@ # validate user-specified fmha_fwd API list -set(EXAMPLE_FMHA_FWD_KNOWN_APIS "fwd;fwd_splitkv") -set(EXAMPLE_FMHA_FWD_ENABLE_APIS "fwd" CACHE STRING - "semicolon-separated list of APIs to generate (${EXAMPLE_FMHA_FWD_KNOWN_APIS}) & link, or \"all\".") -if(EXAMPLE_FMHA_FWD_ENABLE_APIS STREQUAL "all") - set(EXAMPLE_FMHA_FWD_ENABLE_APIS ${EXAMPLE_FMHA_FWD_KNOWN_APIS}) +set(FMHA_FWD_KNOWN_APIS "fwd;fwd_splitkv") +set(FMHA_FWD_ENABLE_APIS "fwd" CACHE STRING + "semicolon-separated list of APIs to generate (${FMHA_FWD_KNOWN_APIS}) & link, or \"all\".") +if(FMHA_FWD_ENABLE_APIS STREQUAL "all") + set(FMHA_FWD_ENABLE_APIS ${FMHA_FWD_KNOWN_APIS}) endif() -foreach(api ${EXAMPLE_FMHA_FWD_ENABLE_APIS}) - if(NOT "${api}" IN_LIST EXAMPLE_FMHA_FWD_KNOWN_APIS) - message(FATAL_ERROR "${api} isn't a known api: ${EXAMPLE_FMHA_FWD_KNOWN_APIS}.") +foreach(api ${FMHA_FWD_ENABLE_APIS}) + if(NOT "${api}" IN_LIST FMHA_FWD_KNOWN_APIS) + message(FATAL_ERROR "${api} isn't a known api: ${FMHA_FWD_KNOWN_APIS}.") endif() endforeach() # "fwd" is a must-have api for the fmha_fwd example, add it if not specified -if(NOT "fwd" IN_LIST EXAMPLE_FMHA_FWD_ENABLE_APIS) - list(APPEND EXAMPLE_FMHA_FWD_ENABLE_APIS "fwd") +if(NOT "fwd" IN_LIST FMHA_FWD_ENABLE_APIS) + list(APPEND FMHA_FWD_ENABLE_APIS "fwd") endif() -string(REPLACE ";" "," EXAMPLE_FMHA_FWD_APIS "${EXAMPLE_FMHA_FWD_ENABLE_APIS}") +string(REPLACE ";" "," FMHA_FWD_APIS "${FMHA_FWD_ENABLE_APIS}") # generate a list of kernels, but not actually emit files at config sta execute_process( COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/generate.py - --api ${EXAMPLE_FMHA_FWD_APIS} --list_blobs ${CMAKE_CURRENT_BINARY_DIR}/fwd_blob_list.txt + --api ${FMHA_FWD_APIS} --list_blobs ${CMAKE_CURRENT_BINARY_DIR}/fwd_blob_list.txt ) execute_process( @@ -37,7 +37,7 @@ file(STRINGS ${CMAKE_CURRENT_BINARY_DIR}/bwd_blob_list.txt FMHA_BWD_GEN_BLOBS) add_custom_command( OUTPUT ${FMHA_FWD_GEN_BLOBS} COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/generate.py - --api ${EXAMPLE_FMHA_FWD_APIS} --output_dir ${CMAKE_CURRENT_BINARY_DIR} + --api ${FMHA_FWD_APIS} --output_dir ${CMAKE_CURRENT_BINARY_DIR} ) add_custom_command( @@ -82,7 +82,7 @@ else() endif() # conditionally enable call to the fwd_splitkv API in fmha_fwd example -if ("fwd_splitkv" IN_LIST EXAMPLE_FMHA_FWD_ENABLE_APIS) +if ("fwd_splitkv" IN_LIST FMHA_FWD_ENABLE_APIS) list(APPEND EXAMPLE_FMHA_FWD_COMPILE_OPTIONS -DCK_TILE_FMHA_FWD_SPLITKV_API=1) else() list(APPEND EXAMPLE_FMHA_FWD_COMPILE_OPTIONS -DCK_TILE_FMHA_FWD_SPLITKV_API=0) From aba46cd655926a37ed0c897505bb4c3f3f744648 Mon Sep 17 00:00:00 2001 From: "PoYen, Chen" Date: Mon, 8 Jul 2024 09:39:15 +0000 Subject: [PATCH 2/2] Regsiter API handlers automatically --- example/ck_tile/01_fmha/generate.py | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/example/ck_tile/01_fmha/generate.py b/example/ck_tile/01_fmha/generate.py index 27347b4476..375884d327 100644 --- a/example/ck_tile/01_fmha/generate.py +++ b/example/ck_tile/01_fmha/generate.py @@ -5,25 +5,29 @@ import argparse from enum import IntEnum from pathlib import Path +import pkgutil +import sys from typing import List, Optional +import codegen.ops from codegen.cmake_config import * -from codegen.ops import ( - fmha_fwd, - fmha_fwd_splitkv, - fmha_bwd -) - class HandlerId(IntEnum): LIST_BLOBS = 0 WRITE_BLOBS = 1 -handlers = { - 'fwd' : (fmha_fwd.list_blobs, fmha_fwd.write_blobs), - 'fwd_splitkv' : (fmha_fwd_splitkv.list_blobs, fmha_fwd_splitkv.write_blobs), - 'bwd' : (fmha_bwd.list_blobs, fmha_bwd.write_blobs), -} +# 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) + if full_module_name not in sys.modules: + ops.append(importer.find_module(module_name).load_module(module_name)) +unwanted_prefix = 'fmha_' +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(output_dir: Optional[str], api_list : List[str], kernel_filter : Optional[str], receipt, mask_impl) -> None: if output_dir is None: