diff --git a/aiter/jit/core.py b/aiter/jit/core.py index 3e92d9fd34..0c0e8a49d6 100644 --- a/aiter/jit/core.py +++ b/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" diff --git a/aiter/jit/utils/chip_info.py b/aiter/jit/utils/chip_info.py index b91c449e65..fb1d298fec 100644 --- a/aiter/jit/utils/chip_info.py +++ b/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() diff --git a/hsa/codegen.py b/hsa/codegen.py index 1b9a890d90..205866d3f7 100644 --- a/hsa/codegen.py +++ b/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"