diff --git a/pyproject.toml b/pyproject.toml index c137718108..99bce2907f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,7 @@ web = [ # Full runtime dependency set installed by the packaged install.sh. runtime = [ "PyYAML>=6.0", + "itanium-demangler>=1.0", "hyperloom-inference_optimizer[llm,web]", ] test = [ diff --git a/src/hyperloom/agents/kernel/tests/test_bypass_source_resolver.py b/src/hyperloom/agents/kernel/tests/test_bypass_source_resolver.py index 1f0874ef83..b231e6dde3 100644 --- a/src/hyperloom/agents/kernel/tests/test_bypass_source_resolver.py +++ b/src/hyperloom/agents/kernel/tests/test_bypass_source_resolver.py @@ -45,8 +45,8 @@ def test_repo_triton_py_is_editable_but_generated_is_not(): def test_resolve_source_without_symbol_is_unresolved(): # The finder is symbol-driven: no device kernel name -> nothing to look up. - assert resolver.resolve_source("_C::silu_and_mul", framework="vllm") == ("", "unresolved") - assert resolver.resolve_source("", device_kernel_name="") == ("", "unresolved") + assert resolver.resolve_source("_C::silu_and_mul", framework="vllm") == ("", "unresolved", "") + assert resolver.resolve_source("", device_kernel_name="") == ("", "unresolved", "") def test_resolve_source_delegates_to_active_finder(monkeypatch): @@ -54,20 +54,21 @@ def test_resolve_source_delegates_to_active_finder(monkeypatch): def fake_resolve_source(op_name, *, framework="", device_kernel_name=""): calls["args"] = (op_name, framework, device_kernel_name) - return "/opt/vllm/csrc/act.cu", "symbol_index" + return "/opt/vllm/csrc/act.cu", "symbol_index", "" monkeypatch.setattr(source_resolver, "resolve_source", fake_resolve_source) - src, method = resolver.resolve_source( + src, method, reason = resolver.resolve_source( "_C::silu_and_mul", framework="vllm", device_kernel_name="act_kernel" ) assert src == "/opt/vllm/csrc/act.cu" assert method == "symbol_index" + assert reason == "" assert calls["args"] == ("_C::silu_and_mul", "vllm", "act_kernel") def test_resolve_source_finder_miss_is_unresolved(monkeypatch): - monkeypatch.setattr(source_resolver, "resolve_source", lambda *a, **k: ("", "unresolved")) - assert resolver.resolve_source("op::x", device_kernel_name="zzz") == ("", "unresolved") + monkeypatch.setattr(source_resolver, "resolve_source", lambda *a, **k: ("", "unresolved", "")) + assert resolver.resolve_source("op::x", device_kernel_name="zzz") == ("", "unresolved", "") def test_resolve_source_swallows_finder_errors(monkeypatch): @@ -75,7 +76,7 @@ def boom(*a, **k): raise OSError("index build failed") monkeypatch.setattr(source_resolver, "resolve_source", boom) - assert resolver.resolve_source("op::x", device_kernel_name="k") == ("", "unresolved") + assert resolver.resolve_source("op::x", device_kernel_name="k") == ("", "unresolved", "") # --- Triton .py resolution (trace kernel_file + AST def-line pinning) ---------- @@ -127,7 +128,7 @@ def test_resolve_triton_py_launcher_form_parsed(repo_dir): def test_resolve_triton_py_rejects_generated(): src, line, method = resolver.resolve_triton_py("/tmp/torchinductor_x/c.py") - assert (src, line, method) == ("", None, "unresolved") + assert (src, line, method) == ("", None, "non_patchable") def test_triton_def_line_single_unambiguous(repo_dir): diff --git a/src/hyperloom/agents/kernel/tests/test_source_resolver.py b/src/hyperloom/agents/kernel/tests/test_source_resolver.py index da78f2dd53..6a3efa4138 100644 --- a/src/hyperloom/agents/kernel/tests/test_source_resolver.py +++ b/src/hyperloom/agents/kernel/tests/test_source_resolver.py @@ -130,9 +130,9 @@ def test_ck_marker_is_namespace_boundary_not_substring(): def test_ck_detection_falls_back_to_mangled_without_cxxfilt(monkeypatch): - # With c++filt unavailable, a mangled CK symbol must still be classified from + # With demangling unavailable, a mangled CK symbol must still be classified from # its length-prefixed namespace, so the verdict does not depend on binutils. - monkeypatch.setattr(source_resolver, "_cxxfilt_base", lambda _m: "") + monkeypatch.setattr(source_resolver, "_demangle", lambda _m: "") assert source_resolver._non_patchable_kind("_ZN2ck15kernel_moe_gemmIiEEvPf") == "aiter_ck" # A non-CK mangled symbol is not misclassified by the fallback. assert source_resolver._non_patchable_kind("_ZN4vllm11some_kernelIiEEvPf") == "" @@ -146,7 +146,7 @@ def test_legacy_tuple_preserves_non_patchable(tmp_path): ) # non_patchable carries an empty source_file but must NOT collapse to # "unresolved" -- callers distinguish "known not rewritable" from "not found". - assert res.as_legacy_tuple() == ("", "non_patchable") + assert res.as_legacy_tuple() == ("", "non_patchable", "aiter_ck") def test_header_declaration_not_selected_over_definition(tmp_path): @@ -203,9 +203,10 @@ def test_legacy_tuple_shape(tmp_path): res = source_resolver.resolve( "op", framework="vllm", device_kernel_name="void vllm::my_test_kernel()", index=index ) - src, method = res.as_legacy_tuple() + src, method, reason = res.as_legacy_tuple() assert src.endswith("activation_kernels.cu") assert method == "symbol_index" + assert reason == "" # --- env fingerprint ------------------------------------------------------- diff --git a/src/hyperloom/agents/kernel/tools/_bypass_report.py b/src/hyperloom/agents/kernel/tools/_bypass_report.py index 9d9c9b11f0..4247467d7c 100644 --- a/src/hyperloom/agents/kernel/tools/_bypass_report.py +++ b/src/hyperloom/agents/kernel/tools/_bypass_report.py @@ -454,8 +454,8 @@ def build_candidates( finder_patchable: bool | None = None finder_status = "" finder_reason = "" - if not source_file and kname: - source_file, method = resolve_source(op_name, framework=framework, device_kernel_name=kname) + if not source_file and kname and source_method == "unresolved": + source_file, method, reason = resolve_source(op_name, framework=framework, device_kernel_name=kname) if source_file: source_method = method finder_patchable = True @@ -468,12 +468,11 @@ def build_candidates( source_method = method finder_patchable = False finder_status = "non_rewritable" - finder_reason = ( - "non-patchable kernel (symbol-detected: CK template / no single editable __global__ source)" - ) - # Repo-scan is the last resort, and only for a genuine miss -- never when - # the finder already returned an authoritative non_patchable verdict. - if not source_file and kname and source_method != "non_patchable": + finder_reason = f"non-patchable kernel ({reason})" if reason else "non-patchable kernel" + # Repo-scan is the last resort, and only for a genuine miss -- never + # when the finder already returned an authoritative non_patchable + # verdict. + if not source_file and kname and source_method == "unresolved": source_file, method = resolve_by_kernel_name(kname) if source_file: source_method = method @@ -563,7 +562,15 @@ def build_candidates( else ( "" if source_file - else (f"source: {finder_reason}" if finder_patchable is False else "source file not resolved") + else ( + f"source: {finder_reason}" + if finder_patchable is False + else ( + "source: non-patchable kernel (trace kernel_file is not editable)" + if source_method == "non_patchable" + else "source file not resolved" + ) + ) ) ), "recommended_backends": list(_REUSABLE_BACKENDS) if kc.reusable else [], diff --git a/src/hyperloom/agents/kernel/tools/_bypass_source_resolver.py b/src/hyperloom/agents/kernel/tools/_bypass_source_resolver.py index 09012d353a..a1d9c79da4 100644 --- a/src/hyperloom/agents/kernel/tools/_bypass_source_resolver.py +++ b/src/hyperloom/agents/kernel/tools/_bypass_source_resolver.py @@ -90,27 +90,22 @@ def resolve_source( *, framework: str = "", device_kernel_name: str = "", -) -> tuple[str, str]: +) -> tuple[str, str, str]: """Resolve a native kernel to its live installed source via the active finder. This is the deterministic op->source tier for the bypass route: it delegates to :func:`source_resolver.resolve_source`, which demangles the device kernel symbol and looks it up in a live ``__global__`` index (method ``"symbol_index"``). There is no static op_to_source map; any import/lookup - failure yields ``("", "unresolved")`` so the caller can fall back to the + failure yields ``("", "unresolved", "")`` so the caller can fall back to the repo-scan tier. - Args: - op_name: The launching op name (carried for reporting, not lookup). - framework: Serving framework hint used to rank multi-tree matches. - device_kernel_name: Device kernel symbol from the trace (authoritative). - Returns: - ``(source_file, "symbol_index")`` on a hit, else ``("", "unresolved")`` - / ``("", "non_patchable")``. + ``(source_file, method, reason)`` — reason carries the specific + non-patchable kind (e.g. ``"tensile_precompiled"``). """ if not device_kernel_name: - return "", "unresolved" + return "", "unresolved", "" try: try: # package import (TraceLens route / tests) from . import source_resolver @@ -120,7 +115,7 @@ def resolve_source( return source_resolver.resolve_source(op_name, framework=framework, device_kernel_name=device_kernel_name) except (ImportError, OSError, ValueError) as exc: log.debug("bypass resolve_source failed for %r: %s", device_kernel_name, exc) - return "", "unresolved" + return "", "unresolved", "" # Triton kernel definition: @triton.jit then optional decorators then def NAME. @@ -133,7 +128,7 @@ def resolve_source( # Directories/paths to skip while scanning source repos. _SCAN_SKIP_MARKERS = ("/__pycache__", "/3rdparty/", "/example", "/test", "/jit/build/", "/.git/") _TRITON_SCAN_EXTS = (".py",) -_NATIVE_SCAN_EXTS = (".cu", ".cuh", ".hip", ".h") +_NATIVE_SCAN_EXTS = (".cu", ".cuh", ".hip", ".h", ".hpp") def _demangle_kernel_name(name: str) -> str | None: @@ -401,14 +396,15 @@ def resolve_triton_py( and the exact ``@triton.jit`` def line is pinned via :func:`triton_def_line`. The AST step is a pure refinement: a resolved file is returned even when the def line cannot be pinned. ``method`` is ``"trace_kernel_file_ast"`` (path + - pinned line), ``"trace_kernel_file"`` (path only), or ``"unresolved"``. + pinned line), ``"trace_kernel_file"`` (path only), ``"non_patchable"``, + or ``"unresolved"``. """ path, line, func = _parse_launcher_form(kernel_file) if not path: return "", None, "unresolved" source = editable_trace_source(path, kernel_kind) if not source: - return "", None, "unresolved" + return "", None, "non_patchable" ast_line: int | None = None if source.lower().endswith(".py") and os.path.isfile(source): ast_line = triton_def_line(source, func=func, symbol=symbol) diff --git a/src/hyperloom/agents/kernel/tools/source_resolver.py b/src/hyperloom/agents/kernel/tools/source_resolver.py index fbe401e2c7..2ff4aacea5 100644 --- a/src/hyperloom/agents/kernel/tools/source_resolver.py +++ b/src/hyperloom/agents/kernel/tools/source_resolver.py @@ -45,6 +45,15 @@ log = logging.getLogger(__name__) +try: + from itanium_demangler import parse as _itanium_parse +except ImportError: + _itanium_parse = None + log.warning( + "itanium-demangler is not installed. Kernel classification may be degraded. " + "Install it with: pip install itanium-demangler" + ) + try: # package import (TraceLens route / tests) from . import kernel_source_index, source_env from ._bypass_source_resolver import is_editable_source @@ -100,17 +109,18 @@ class ResolveResult: elapsed_ms: float reason: str = "" - def as_legacy_tuple(self) -> tuple[str, str]: - """Legacy ``(source_file, method)`` shape for drop-in compatibility. + def as_legacy_tuple(self) -> tuple[str, str, str]: + """Legacy ``(source_file, method, reason)`` shape. A hit keeps its ``method`` (``"symbol_index"``); ``"non_patchable"`` is preserved even though its ``source_file`` is empty (so callers can tell "known not rewritable" from "not found"); every other empty-source - outcome collapses to ``"unresolved"``. + outcome collapses to ``"unresolved"``. ``reason`` carries the specific + non-patchable kind (e.g. ``"tensile_precompiled"``). """ if self.source_file or self.method == "non_patchable": - return (self.source_file, self.method) - return ("", "unresolved") + return (self.source_file, self.method, self.reason) + return ("", "unresolved", "") # ---------------------------------------------------------------------------- @@ -177,8 +187,9 @@ def _pct(p: float) -> float: def _cxxfilt_base(mangled: str) -> str: """Demangle via ``c++filt`` when available (``""`` on failure). - Cached: demangling is pure and the same mangled symbols recur across - candidates, so we pay the subprocess spawn at most once per symbol. + Fallback for when ``itanium-demangler`` is not installed. Cached: the same + mangled symbols recur across candidates, so subprocess spawn cost is paid at + most once per symbol. """ if not shutil.which("c++filt"): return "" @@ -189,19 +200,52 @@ def _cxxfilt_base(mangled: str) -> str: text=True, timeout=5, ) - return proc.stdout.strip() + result = proc.stdout.strip() + # c++filt returns the input unchanged on failure — treat that as no result. + return result if result != mangled else "" except (OSError, subprocess.SubprocessError) as exc: log.debug("c++filt demangle failed for %r: %s", mangled, exc) return "" +@functools.lru_cache(maxsize=8192) +def _demangle(mangled: str) -> str: + """Demangle an Itanium-mangled symbol via ``itanium-demangler``. + + Returns the demangled string on success, ``""`` on parse failure. Unlike + ``c++filt``, this is pure-Python (no subprocess) and handles deeply-nested + CK template instantiations that exceed ``c++filt``'s recursion limits. + + Falls back to ``_cxxfilt_base`` when ``itanium-demangler`` is not installed. + """ + if _itanium_parse is None: + return _cxxfilt_base(mangled) + try: + node = _itanium_parse(mangled) + return str(node) if node is not None else "" + except Exception as exc: # noqa: BLE001 — malformed symbols should not propagate. + log.debug("itanium demangle failed for %r: %s", mangled, exc) + return "" + + def _base_from_demangled(name: str) -> str: - """Extract the base kernel identifier from a demangled/plain symbol.""" - # Keep only the head before params/templates, drop namespaces, then take the - # last token (drops any leading return type/qualifiers: "void ns::foo" -> "foo"). - head = re.split(r"[(<]", name.strip(), maxsplit=1)[0].split("::")[-1] - tokens = head.split() - return tokens[-1] if tokens else "" + """Extract the base kernel identifier from a demangled/plain symbol. + + Strips ``void`` return type, ``(anonymous namespace)::`` qualifiers, + template args ``<...>``, function args ``(...)``, and namespace prefixes. + """ + n = (name or "").strip() + if not n: + return "" + if n.startswith("void "): + n = n[len("void "):].strip() + n = n.replace("(anonymous namespace)::", "") + n = re.sub(r"<.*$", "", n) + n = re.sub(r"\(.*$", "", n) + n = n.strip() + if "::" in n: + n = n.rsplit("::", 1)[-1] + return n def _base_from_mangled(mangled: str) -> str: @@ -251,7 +295,7 @@ def base_symbol(device_kernel_name: str) -> str: if not raw: return "" if raw.startswith("_Z"): - demangled = _cxxfilt_base(raw) + demangled = _demangle(raw) if demangled and demangled != raw: return _base_from_demangled(demangled) return _base_from_mangled(raw) @@ -261,24 +305,35 @@ def base_symbol(device_kernel_name: str) -> str: # ---------------------------------------------------------------------------- # Non-patchable detection (symbol-derived, no external metadata) # ---------------------------------------------------------------------------- -def _non_patchable_kind(device_kernel_name: str) -> str: - """Return a non-patchable kind label from the symbol alone (``""`` if none). - - CK (Composable Kernel) template instantiations have no single editable - ``__global__`` source, so they are detected from the symbol's namespace and - reported as ``"aiter_ck"`` (else ``""``). The match is boundary-anchored so - unrelated names ending in ``ck`` are not misclassified, and it falls back to - the mangled form when ``c++filt`` is unavailable so the verdict does not - depend on binutils being installed. +def _non_patchable_kind(device_kernel_name: str, *, op_name: str = "") -> str: + """Return a non-patchable kind label from the symbol/op name (``""`` if none). + + Detected categories: + + * **Tensile** (``Cijk_*``): pre-compiled GPU assembly shipped as ``.co`` + code objects in hipBLASLt/rocBLAS — no ``.cu`` source exists. + * **MIOpen** (``aten::miopen_*`` op name): assembly-generated convolution + kernels pre-compiled into MIOpen's kernel database. Identified by op name + rather than device kernel name patterns. + * **CK** (Composable Kernel, ``ck::`` / ``ck_tile::``): template + instantiations with no single editable ``__global__`` source. """ raw = (device_kernel_name or "").strip() if not raw: return "" + # Tensile GEMM kernels: Cijk___ + # These are pre-compiled assembly (.co), no .cu source exists. + if raw.startswith("Cijk_"): + return "tensile_precompiled" + # MIOpen convolution kernels: identified by op name rather than device + # kernel name patterns, which vary across MIOpen versions. + if "miopen" in (op_name or "").lower(): + return "miopen_precompiled" if raw.startswith("_Z"): - demangled = _cxxfilt_base(raw) + demangled = _demangle(raw) if demangled: return "aiter_ck" if _CK_DEMANGLED_RE.search(demangled.lower()) else "" - # c++filt absent: classify from the mangled namespace prefix instead. + # Demangling failed: classify from the mangled namespace prefix instead. return "aiter_ck" if _CK_MANGLED_RE.search(raw) else "" return "aiter_ck" if _CK_DEMANGLED_RE.search(raw.lower()) else "" @@ -348,7 +403,7 @@ def finish(res: ResolveResult) -> ResolveResult: # Cheap gate first (symbol-derived): CK template instantiations have no # single editable source, so bail with a clear reason. - nonp_kind = _non_patchable_kind(device_kernel_name) + nonp_kind = _non_patchable_kind(device_kernel_name, op_name=op_name) if nonp_kind: return finish( ResolveResult(