fix_prebuiild_asm

aiter/jit/core.py

@@ -20,7 +20,7 @@
 
 this_dir = os.path.dirname(os.path.abspath(__file__))
 sys.path.insert(0, f"{this_dir}/utils/")
-from chip_info import get_gfx
+from chip_info import get_gfx, get_gfx_list
 from cpp_extension import _jit_compile, get_hip_version
 from file_baton import FileBaton
 from torch_guard import torch_compile_guard  # noqa: E402
@@ -257,9 +257,16 @@ def get_config_file(env_name, default_file, tuned_file_name):
 sys.path.insert(0, AITER_META_DIR)
 AITER_CSRC_DIR = f"{AITER_META_DIR}/csrc"
 AITER_GRADLIB_DIR = f"{AITER_META_DIR}/gradlib"
-gfx = get_gfx()
-AITER_ASM_DIR = f"{AITER_META_DIR}/hsa/{gfx}/"
-os.environ["AITER_ASM_DIR"] = AITER_ASM_DIR
+gfx = get_gfx_list()
+if len(gfx) == 1:
+    # single GPU arch
+    AITER_ASM_DIR = f"{AITER_META_DIR}/hsa/{gfx[0]}/"
+    os.environ["AITER_ASM_DIR"] = AITER_ASM_DIR
+else:
+    # multiple GPU archs
+    AITER_ASM_DIR = [f"{AITER_META_DIR}/hsa/{g}/" for g in gfx]
+    os.environ["AITER_ASM_DIR"] = ":".join(AITER_ASM_DIR)
+
 CK_3RDPARTY_DIR = os.environ.get(
     "CK_DIR", f"{AITER_META_DIR}/3rdparty/composable_kernel"
 )

aiter/jit/optCompilerConfig.json

@@ -310,7 +310,7 @@
         "extra_ldflags": "None",
         "extra_include": [],
         "verbose": "False",
-        "blob_gen_cmd": "f'{get_asm_dir()}/bf16gemm/codegen.py --output_dir {{}}'"
+        "blob_gen_cmd": "f'{AITER_META_DIR}/hsa/codegen.py -m bf16gemm --output_dir {{}}'"
     },
     "module_gemm_a4w4_asm": {
         "srcs": [

csrc/py_itfs_cu/asm_gemm_a16w16.cu

@@ -62,6 +62,8 @@ get_heuristic_kernel(int M,
     hipDeviceProp_t dev_prop;
     HIP_CALL(hipGetDevice(&dev));
     HIP_CALL(hipGetDeviceProperties(&dev_prop, dev));
+    std::string arch_id = get_gpu_arch();
+    printf("Arch ID: %s\n", arch_id.c_str());
     uint32_t num_cu = dev_prop.multiProcessorCount;
     // printf("num_cu: %d\n", num_cu);
     uint32_t empty_cu      = num_cu;
@@ -75,6 +77,8 @@ get_heuristic_kernel(int M,
 
     for(const auto& el : *cfgs)
     {
+        if (el.first.find(arch_id) != 0)
+            continue;
         const auto& cfg = el.second;
         if(kernelName.has_value() && kernelName.value() != el.first)
             continue;
@@ -125,6 +129,7 @@ get_heuristic_kernel(int M,
                 compute2mem_effi   = local_compute2mem_effi;
                 oob                = (M % cfg.tileM == 0) ? 0 : cfg.tileM - (M % cfg.tileM);
                 selectedKernelName = el.first;
+                printf("Selected Kernel: %s\n", selectedKernelName.c_str());
                 selectedsplitK     = split_K;
             }
         }
@@ -243,7 +248,7 @@ torch::Tensor gemm_a16w16_asm(torch::Tensor& A,   // A:[M, K] bf16
     if(it_kl != config_map->end())
     {
         const auto& cfg     = it_kl->second;
-        const char* name    = cfg.name.c_str();
+        const char* name    = cfg.knl_name.c_str();
         const char* co_name = cfg.co_name.c_str();
         SUBM                = cfg.tileM;
         SUBN                = cfg.tileN;

hsa/codegen.py

@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+import os
+import sys
+import argparse
+import glob
+import pandas as pd
+from collections import defaultdict
+
+this_dir = os.path.dirname(os.path.abspath(__file__))
+base_dir = os.path.basename(this_dir)
+archs = [os.path.basename(os.path.normpath(path)) for path in os.environ.get("AITER_ASM_DIR").split(":")]
+
+
+content =  """// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+#pragma once
+#include <unordered_map>
+
+"""
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        prog="generate",
+        description="gen API for asm Bf16_gemm kernel",
+    )
+    parser.add_argument(
+        "-m",
+        "--module",
+        required=True,
+        help="""module of ASM kernle
+    e.g.: -m bf16gemm
+""",
+    )
+    parser.add_argument(
+        "-o",
+        "--output_dir",
+        default="aiter/jit/build",
+        required=False,
+        help="write all the blobs into a directory",
+    )
+    args = parser.parse_args()
+    cfgs = []
+
+    csv_groups = defaultdict(list)
+    for arch in archs:
+        # print(f"{this_dir}/{arch}/{args.module}")
+        for el in glob.glob(f"{this_dir}/{arch}/{args.module}/*.csv", recursive=True):
+            df = pd.read_csv(el)
+            cfgname = os.path.basename(el).split(".")[0]
+            csv_groups[cfgname].append({
+                'file_path': el,
+                'arch': arch
+            })
+
+    ## deal with same name csv
+    cfgs = []
+    have_get_header = False
+    for cfgname, file_info_list in csv_groups.items():
+        dfs = []
+        for file_info in file_info_list:
+            single_file = file_info["file_path"]
+            arch = file_info['arch']
+            df = pd.read_csv(single_file)
+            # check headers
+            headers_list = df.columns.tolist()
+            required_columns = {'knl_name', 'co_name'}
+            if not required_columns.issubset(headers_list):
+                missing = required_columns - set(headers_list)
+                print(f"ERROR: Invalid assembly CSV format -- {single_file}. Missing required columns: {', '.join(missing)}")
+            df['arch'] = arch    # add arch into df
+            dfs.append(df)
+        if dfs:
+            combine_df = pd.concat(dfs, ignore_index=True).fillna(0)
+            if have_get_header == False:
+                headers_list = combine_df.columns.tolist()
+                required_columns = {'knl_name', 'co_name', 'arch'}
+                other_columns = [col for col in headers_list if col not in required_columns]
+                other_columns_comma = ', '.join(other_columns)
+                other_columns_cpp_def =  "\n".join([f"    int {col};" for col in other_columns])
+                content += f'''
+#define ADD_CFG({other_columns_comma}, path, knl_name, co_name, arch)         \\
+    {{                                         \\
+        arch knl_name, {{ knl_name, path co_name, arch, {other_columns_comma} }}         \\
+    }}
+
+struct {args.module}Config
+{{
+    std::string knl_name;
+    std::string co_name;
+    std::string arch;
+{other_columns_cpp_def}
+}};
+
+using CFG = std::unordered_map<std::string, {args.module}Config>;
+
+'''
+                have_get_header = True
+            cfg = [
+                f'ADD_CFG({", ".join(str(getattr(row, col)) for col in other_columns)}, "{os.path.relpath(os.path.dirname(el), f"{this_dir}/{arch}")}/", "{row.knl_name}", "{row.co_name}", "{row.arch}"),'
+                for row in combine_df.itertuples(index=False)
+            ]
+            cfg_txt = "\n            ".join(cfg) + "\n"
+
+            txt = f"""static CFG cfg_{cfgname} = {{
+            {cfg_txt}}};"""
+            cfgs.append(txt)
+
+        content += "\n".join(cfgs) + "\n"
+
+    with open(f"{args.output_dir}/asm_{args.module}_configs.hpp", "w") as f:
+        f.write(content)

hsa/gfx942/bf16gemm/bf16gemm_outf32.csv

@@ -1,4 +1,4 @@
-knl_name,knl_file,tn,tileM,tileN,pf
+knl_name,co_name,tn,tileM,tileN,pf
 _ZN5aiter28bf16gemm_outf32_tn_32x64_pf3E,bf16gemm_outf32_tn_32x64_pf3.co,1,32,64,3
 _ZN5aiter28bf16gemm_outf32_tn_48x64_pf3E,bf16gemm_outf32_tn_48x64_pf3.co,1,48,64,3
 _ZN5aiter28bf16gemm_outf32_tn_64x64_pf3E,bf16gemm_outf32_tn_64x64_pf3.co,1,64,64,3

hsa/gfx950/bf16gemm/bf16gemm_outf32.csv

@@ -1,4 +1,4 @@
-knl_name,knl_file,tn,tileM,tileN,pf
+knl_name,co_name,tn,tileM,tileN,pf
 _ZN5aiter28bf16gemm_outf32_tn_32x64_pf3E,bf16gemm_outf32_tn_32x64_pf3.co,1,32,64,3
 _ZN5aiter28bf16gemm_outf32_tn_48x64_pf3E,bf16gemm_outf32_tn_48x64_pf3.co,1,48,64,3
 _ZN5aiter28bf16gemm_outf32_tn_64x64_pf3E,bf16gemm_outf32_tn_64x64_pf3.co,1,64,64,3