[hiptools] robust kernel demangle parsing and double kernarg support

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,19 +105,66 @@ 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()
 
 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")

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"])