update codegen

aiter/jit/core.py

@@ -265,15 +265,9 @@ 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_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)
+gfxs = get_gfx_list()
+AITER_ASM_DIR = f"{AITER_META_DIR}/hsa/{get_gfx()}/"
+os.environ["AITER_ASM_DIR"] = AITER_ASM_DIR
 
 CK_3RDPARTY_DIR = os.environ.get(
     "CK_DIR", f"{AITER_META_DIR}/3rdparty/composable_kernel"

aiter/jit/utils/chip_info.py

@@ -4,9 +4,9 @@
 import os
 import re
 import subprocess
-from torch_guard import torch_compile_guard
 
 from cpp_extension import executable_path
+from torch_guard import torch_compile_guard
 
 GFX_MAP = {
     0: "native",
@@ -91,9 +91,15 @@ def get_gfx_list() -> list[str]:
 
     gfx_env = os.getenv("GPU_ARCHS", "native")
     if gfx_env == "native":
-        return _detect_native()
+        try:
+            gfxs = _detect_native()
+        except RuntimeError:
+            gfxs = ["cpu"]
+    else:
+        gfxs = [g.strip() for g in gfx_env.split(";") if g.strip()]
+    os.environ["AITER_GPU_ARCHS"] = ";".join(gfxs)
 
-    return [g.strip() for g in gfx_env.split(";") if g.strip()]
+    return gfxs
 
 
 @torch_compile_guard()

hsa/codegen.py

@@ -1,19 +1,20 @@
 # 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
-import numpy as np
+import os
+import sys
 from collections import defaultdict
 
+import numpy as np
+import pandas as pd
+
 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(":")
+archs = [el for el in os.environ["AITER_GPU_ARCHS"].split(";")]
+archs_supported = [
+    os.path.basename(os.path.normpath(path)) for path in glob.glob(f"{this_dir}/*/")
 ]
 
 
@@ -48,12 +49,10 @@
     cfgs = []
 
     csv_groups = defaultdict(list)
-    for arch in archs:
-        # print(f"{this_dir}/{arch}/{args.module}")
+    for arch in archs_supported:
         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})
 
@@ -115,9 +114,18 @@
 """
                 have_get_header = True
             cfg = [
-                f'ADD_CFG({", ".join(f"\"{getattr(row, col)}\"" if not str(getattr(row, col)).isdigit() else f"{getattr(row, col):>4}" for col in other_columns)}, '
-                f'"{row.arch}", "{relpath}/", "{row.knl_name}", "{row.co_name}"),'
+                "ADD_CFG("
+                + ", ".join(
+                    (
+                        f"{int(getattr(row, col)):>4}"
+                        if str(getattr(row, col)).replace(".", "", 1).isdigit()
+                        else f'"{getattr(row, col)}"'
+                    )
+                    for col in other_columns
+                )
+                + f', "{row.arch}", "{relpath}/", "{row.knl_name}", "{row.co_name}"),'
                 for row in combine_df.itertuples(index=False)
+                if row.arch in archs
             ]
             cfg_txt = "\n    ".join(cfg) + "\n"