Add xDiT kernel source mapping support

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 = [

src/hyperloom/agents/kernel/tests/test_bypass_source_resolver.py

@@ -45,37 +45,38 @@ 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):
     calls = {}
 
     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):
     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):

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<float>()", 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 -------------------------------------------------------

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 [],

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)

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_<A_layout>_<B_layout>_<config...>
+    # 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(