From 95bffe3d45675d535cfdf60a6dcb74449f58d08b Mon Sep 17 00:00:00 2001 From: apinge Date: Tue, 14 Apr 2026 11:22:39 +0000 Subject: [PATCH] first impl Signed-off-by: apinge --- src/core/hiptools.py | 122 ++++++++++++++++++++++++++++++------ tests/core/test_hiptools.py | 34 ++++++++++ 2 files changed, 138 insertions(+), 18 deletions(-) create mode 100644 tests/core/test_hiptools.py diff --git a/src/core/hiptools.py b/src/core/hiptools.py index 8ea49cbc..336ee1c5 100644 --- a/src/core/hiptools.py +++ b/src/core/hiptools.py @@ -1,11 +1,13 @@ import ctypes from ctypes.util import find_library +import re import subprocess import functools import os, sys import inspect from .asmtools import prettify import torch +from typing import List, Optional, Tuple @functools.cache def get_lib(): @@ -69,6 +71,8 @@ def __init__(self, module_fpath, sym_name, kname, kargs): fields.append((f"arg_{i}", ctypes.c_uint)) elif arg_type == "float" or arg_type == "float32_t": fields.append((f"arg_{i}", ctypes.c_float)) + elif arg_type == "double": + fields.append((f"arg_{i}", ctypes.c_double)) else: raise Exception(f"Unsupported arg type: {arg_type}") class Args(ctypes.Structure): @@ -101,7 +105,12 @@ def __call__(self, gridDims:list[int], blockDims:list[int], *args, sharedMemByte @functools.cache def amdgpu_arch(): - gfx_archs = subprocess.check_output(["/opt/rocm/llvm/bin/amdgpu-arch"]).decode('utf-8') + # 优先用环境变量指定 ISA；否则调用 amdgpu-arch 列出本机 GPU + # 方便交叉编译 + override = os.environ.get("PYHIP_AMDGPU_ARCH", "").strip() + if override: + return override + gfx_archs = subprocess.check_output(["/opt/rocm/llvm/bin/amdgpu-arch"]).decode("utf-8") # index = torch.cuda.current_device() index = 0 return gfx_archs.splitlines()[index].strip() @@ -109,11 +118,53 @@ def amdgpu_arch(): PYHIP_CACHE_DIR = os.getenv("PYHIP_CACHE_DIR", os.path.expanduser("~/.pyhip")) os.makedirs(PYHIP_CACHE_DIR, exist_ok=True) +# 匹配 #include "relative/path.hpp"（不含系统头 <...>） +_RE_INCLUDE_QUOTED = re.compile(r'^\s*#\s*include\s+"([^"]+)"') + +def _device_src_mtime_with_local_includes(src_path: str) -> float: + """ + 主源文件 **最后修改时间**（POSIX 里常记作 mtime，即 ``os.path.getmtime``）与其中 + ``#include "..."`` 所引用且存在于磁盘上的文件的 mtime 的**最大值**。 + 路径相对于 **主源文件所在目录** 解析（与编译器 ``-I`` 同目录时行为一致）。 + 不递归展开被包含文件里的 include（避免循环与成本；需要时可改环境变量补依赖）。 + """ + src_path = os.path.abspath(src_path) + mt = os.path.getmtime(src_path) + src_folder = os.path.dirname(src_path) + try: + with open(src_path, encoding="utf-8", errors="ignore") as f: + for line in f: + m = _RE_INCLUDE_QUOTED.match(line) + if not m: + continue + rel = m.group(1).strip() + if not rel or rel.startswith("<"): + continue + inc_path = os.path.normpath(os.path.join(src_folder, rel)) + if os.path.isfile(inc_path): + mt = max(mt, os.path.getmtime(inc_path)) + except OSError: + pass + # 可选：``PYHIP_DEVICE_COMPILE_EXTRA_DEPS=/a/x.hip,/b/y.hip`` 额外纳入 mtime（绝对路径或相对 cwd） + extra = os.environ.get("PYHIP_DEVICE_COMPILE_EXTRA_DEPS", "").strip() + if extra: + for part in extra.split(os.pathsep if os.pathsep in extra else ","): + p = part.strip() + if not p: + continue + p = os.path.abspath(os.path.expanduser(p)) + if os.path.isfile(p): + mt = max(mt, os.path.getmtime(p)) + return mt + + def compile_hip_device_only(src_path, extra_compiler_options, macros=None): - src_mtime = os.path.getmtime(src_path) + src_mtime = _device_src_mtime_with_local_includes(src_path) src_folder, src_filename = os.path.split(src_path) pre, ext = os.path.splitext(src_filename) + # 使与源文件同目录的头文件（如显式实例化）可被 #include "" + inc_dir = f"-I{src_folder}" if src_folder else "-I." # put all intermedite file under cache pre = f"{PYHIP_CACHE_DIR}/{pre}" @@ -124,22 +175,25 @@ def compile_hip_device_only(src_path, extra_compiler_options, macros=None): extra_compiler_options += f" -D{k}={v}" pre += f"-{k}={v}" - assert ext == ".cpp" or ext == ".s" + gfx_arch = amdgpu_arch() + # 不同 offload-arch 必须分文件缓存，否则 .co 会错用其它 ISA 的汇编 + pre = f"{pre}-{gfx_arch}" + + assert ext in (".cpp", ".s", ".hip") ll_path = pre + ".ll" asm_path = pre + ".s" co_path = pre + ".co" - gfx_arch = amdgpu_arch() if os.getenv("DUMP_LL", None) is not None: - cmd1 = f"hipcc -x hip --offload-device-only --offload-arch={gfx_arch} -std=c++20 -I. -O2 {src_path} {extra_compiler_options} -S -emit-llvm -o {ll_path}" + cmd1 = f"hipcc -x hip --offload-device-only --offload-arch={gfx_arch} -std=c++20 {inc_dir} -O2 {src_path} {extra_compiler_options} -S -emit-llvm -o {ll_path}" print(cmd1) if os.system(cmd1) != 0: raise Exception("compilation 0 failed") print(f"\033[0;32m LLVM IR {ll_path} was generated. \033[0m") - if ext == ".cpp" and (not os.path.isfile(asm_path) or src_mtime > os.path.getmtime(asm_path)): + if ext in (".cpp", ".hip") and (not os.path.isfile(asm_path) or src_mtime > os.path.getmtime(asm_path)): # (re)compile into asm - cmd1 = f"hipcc -x hip --offload-device-only --offload-arch={gfx_arch} -std=c++20 -I. -O2 {extra_compiler_options} -Rpass-analysis=kernel-resource-usage {src_path} -S -o {asm_path}" + cmd1 = f"hipcc -x hip --offload-device-only --offload-arch={gfx_arch} -std=c++20 {inc_dir} -O2 {extra_compiler_options} -Rpass-analysis=kernel-resource-usage {src_path} -S -o {asm_path}" print(cmd1) if os.system(cmd1) != 0: raise Exception("compilation 1 failed") @@ -153,23 +207,55 @@ def compile_hip_device_only(src_path, extra_compiler_options, macros=None): raise Exception("compilation 2 failed") return co_path + +def _parse_kernel_demangle_func_sig(func_sig: str) -> Optional[Tuple[str, List[str]]]: + """ + 从 llvm-objdump --demangle 拼出的 ``func_sig``（``join(ls[6:])``）解析出 + ``(fname, [arg types...])``；无法解析（无括号）时返回 None。 + 供 ``get_all_kernel_args`` 与单测共用。 + """ + func_sig = func_sig.strip().rstrip(")") + if "(" not in func_sig: + return None + fname, args_str = func_sig.split("(", 1) + fname = fname.strip() + if fname.startswith("void "): + fname = fname[5:].strip() + args_str = args_str.rstrip(")").strip() + if not args_str: + args = [] + else: + args = [a.strip() for a in args_str.split(",")] + return fname, args + + @functools.cache def get_all_kernel_args(co_path): - # we need both demangle & symbol name for loading & argtype parsing - dynamic_syms = subprocess.check_output(["/opt/rocm/llvm/bin/llvm-objdump", "--dynamic-syms", "--demangle", co_path]).decode('utf-8') - dynamic_syms_raw = subprocess.check_output(["/opt/rocm/llvm/bin/llvm-objdump", "--dynamic-syms", co_path]).decode('utf-8') + """ + 从 .co 的 ELF 动态符号表解析每个 kernel 的「Python 侧函数名」与「参数类型」。 + hipModuleGetFunction 必须用 **未 demangle 的符号名**（raw），而参数类型串只在 + ``--demangle`` 输出里可读，因此跑两次 llvm-objdump。 + """ + dynamic_syms = subprocess.check_output( + ["/opt/rocm/llvm/bin/llvm-objdump", "--dynamic-syms", "--demangle", co_path] + ).decode("utf-8") + dynamic_syms_raw = subprocess.check_output( + ["/opt/rocm/llvm/bin/llvm-objdump", "--dynamic-syms", co_path] + ).decode("utf-8") kernel_args = {} for line, line_raw in zip(dynamic_syms.splitlines(), dynamic_syms_raw.splitlines()): ls = line.split() - if len(ls) < 7: continue - if ls[3] != ".text" : continue + if len(ls) < 7: + continue + if ls[3] != ".text": + continue + # 与 hipModuleLoad 后 GetFunction 传入的字符串一致，必须用 raw 列里的名字（常为 _Z 修饰名） symbol_name = line_raw.split()[6] func_sig = " ".join(ls[6:]).strip().rstrip(")") - fname, args = func_sig.split("(") - if len(args) == 0: - args = [] - else: - args = [a.strip() for a in args.split(",")] + parsed = _parse_kernel_demangle_func_sig(func_sig) + if parsed is None: + continue + fname, args = parsed kernel_args[fname] = (symbol_name, args) return kernel_args @@ -188,7 +274,7 @@ def test_kernel(q, k, v, n, sm): ... @functools.cache def _build_hip_src(src_fpath, extra_compiler_options, macros): - if src_fpath.endswith(".cpp") or src_fpath.endswith(".s"): + if src_fpath.endswith((".cpp", ".hip", ".s")): module_fpath = compile_hip_device_only(src_fpath, extra_compiler_options, macros) else: assert src_fpath.endswith(".co") diff --git a/tests/core/test_hiptools.py b/tests/core/test_hiptools.py new file mode 100644 index 00000000..223a7e8f --- /dev/null +++ b/tests/core/test_hiptools.py @@ -0,0 +1,34 @@ +"""测试 ``hiptools._parse_kernel_demangle_func_sig``（demangle 串 → 函数名与参数类型列表）。 +python -m pytest tests/core/test_hiptools.py -v +""" +from __future__ import annotations + +from pyhip.core.hiptools import _parse_kernel_demangle_func_sig + + +def test_strip_void_prefix_matches_python_func_name(): + """demangle 常带 ``void foo(...)``；pyhip 用 ``getattr(mod, \"foo\")``，必须去掉 ``void ``。""" + assert _parse_kernel_demangle_func_sig("void bar(int, float)") == ("bar", ["int", "float"]) + + +def test_no_paren_returns_none(): + assert _parse_kernel_demangle_func_sig("plain_symbol") is None + + +def test_split_first_open_paren_only(): + """只用第一个 ``(`` 分开函数名与参数；避免无上限 split 把模板里括号拆碎。""" + r = _parse_kernel_demangle_func_sig("void baz(int, float)") + assert r is not None + fname, args = r + assert fname == "baz" + assert args == ["int", "float"] + + +def test_empty_arg_list(): + assert _parse_kernel_demangle_func_sig("void qux()") == ("qux", []) + + +def test_comma_in_types_simple(): + """参数类型里逗号分隔；若类型名含逗号（模板）会误切，属已知限制。""" + r = _parse_kernel_demangle_func_sig("void k(void const*, int)") + assert r == ("k", ["void const*", "int"])