Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -190,6 +190,11 @@ hyperloom = [
"hyperloom.agents.kernel" = [
"scripts/*.sh",
]
# Vendor-operator-playbook registry (see tools/_vendor_operator_playbooks.py):
# read at runtime by classify_patchability()/forge_submit.py, so it must ship.
"hyperloom.agents.kernel.tools" = [
"vendor_operator_playbooks.json",
]
# ``agentx/deploy.py`` copies these into InferenceX ``benchmarks/`` at runtime
# and raises FileNotFoundError when they are absent from the install.
"hyperloom.inference_optimizer.assets.agentx" = [
Expand Down

Large diffs are not rendered by default.

201 changes: 201 additions & 0 deletions src/hyperloom/agents/kernel/tools/_vendor_operator_playbooks.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,201 @@
###############################################################################
# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc.
# SPDX-License-Identifier: MIT
#
# See LICENSE for license information.
###############################################################################

"""Vendor-operator-playbook registry: route a closed-source hot kernel to a
validated KernelForge *task bundle* instead of a source rewrite.

Most of the forge-submission pipeline assumes a hot kernel has an editable
device source file (``kernel_url`` -> in-place rewrite). Some vendor
operators -- mori's EP dispatch/combine all-to-all is the first case -- are
pip-installed compiled libraries with no such source, but do have a small,
named set of launch-config knobs that a KernelForge forge-loop task bundle
has already been validated to tune (see KernelForge PR #88's "Making this
real" section for the design rationale this module implements).

This is deliberately a narrow, explicit carve-out (one JSON registry, sibling
to the retired ``op_to_source.json``) rather than a general "config-tuning"
system: a candidate only gets vendor-playbook treatment when it matches a
registry entry by name.
"""

from __future__ import annotations

import copy
import functools
import json
import os
from pathlib import Path
from typing import Any

_REGISTRY_PATH = Path(__file__).resolve().parent / "vendor_operator_playbooks.json"


@functools.lru_cache(maxsize=1)
def load_vendor_operator_playbooks() -> tuple[dict[str, Any], ...]:
"""Load and cache the vendor-operator-playbook registry.

Returns:
A tuple of playbook entry dicts (empty when the registry file is
missing or malformed -- a missing registry must never be fatal to
the rest of the classification pipeline).
"""
try:
raw = json.loads(_REGISTRY_PATH.read_text(encoding="utf-8"))
except (OSError, ValueError):
return ()
playbooks = raw.get("playbooks") if isinstance(raw, dict) else None
if not isinstance(playbooks, list):
return ()
return tuple(entry for entry in playbooks if isinstance(entry, dict) and entry.get("id"))


def _reset_vendor_operator_playbooks_cache() -> None:
"""Clear the cached registry (tests only, e.g. after monkeypatching the path)."""
load_vendor_operator_playbooks.cache_clear()


def _candidate_haystack(candidate: dict[str, Any]) -> str:
"""Join every text field a playbook's ``any_marker`` may match against."""
fields = (
candidate.get("name"),
candidate.get("operation"),
candidate.get("library"),
candidate.get("source_file"),
candidate.get("kernel_repo"),
)
return " ".join(str(f or "") for f in fields).lower()


def _last_symbol_segment(value: str) -> str:
"""Return the trailing method/function segment of a qualified symbol.

``mori::EpDispatchCombineOp::combine`` -> ``combine``; a plain name with
no separator is returned unchanged. Needed because a class name like
``EpDispatchCombineOp`` itself contains the substring "dispatch", so
matching a role marker against the *whole* qualified name is ambiguous --
only the actual called method disambiguates dispatch vs combine.
"""
tail = value
for sep in ("::", ".", "/"):
tail = tail.rsplit(sep, 1)[-1]
return tail


def _role_haystack(candidate: dict[str, Any]) -> str:
"""Return the field(s) a playbook's ``name_any`` (op-role pattern) should match.

Prefers ``operation`` (often already the specific call, e.g.
``"combine"``) over ``name`` (which may be a fully-qualified
``Class::method`` symbol whose class name can itself contain another
role's marker); either way, only the trailing symbol segment is
matched, never the whole qualified string. This repo's own convention
(``_task_group_contract.logical_operator_name()``,
``_bypass_report.py``'s task-group builder) is to set ``operation`` to
the fully-qualified name too, e.g. ``mori::EpDispatchCombineOp::combine``
-- taking ``operation`` verbatim would silently reintroduce the exact
dispatch/combine ambiguity this function exists to resolve the moment
some producer starts populating that field on a candidate row (PR #1191
review finding #6).
"""
operation = str(candidate.get("operation") or "").strip()
if operation:
return _last_symbol_segment(operation).lower()
return _last_symbol_segment(str(candidate.get("name") or "")).lower()


def match_vendor_operator_playbook(candidate: dict[str, Any]) -> dict[str, Any] | None:
"""Return a matched playbook entry for ``candidate``, or ``None``.

A candidate matches a playbook when at least one of the playbook's
``any_marker`` strings appears somewhere in the candidate's identifying
fields (name/operation/library/source_file/kernel_repo) AND at least one
of its ``name_any`` strings appears in the candidate's name/operation --
e.g. mori's playbook requires both "mori" (library/source evidence) and
"dispatch" or "combine" (which op within mori this is).

Args:
candidate: The hot-kernel candidate dict (as built by
``tracelens_analysis``).

Returns:
A deep copy of the matched registry entry, augmented with a
``"role"`` key set to whichever ``name_any`` marker matched (e.g.
``"dispatch"`` or ``"combine"``), or ``None`` when nothing matches.
"""
if not isinstance(candidate, dict):
return None
haystack = _candidate_haystack(candidate)
role_haystack = _role_haystack(candidate)
if not haystack or not role_haystack:
return None
for playbook in load_vendor_operator_playbooks():
match = playbook.get("match")
if not isinstance(match, dict):
continue
any_markers = [str(m).lower() for m in (match.get("any_marker") or [])]
if any_markers and not any(marker in haystack for marker in any_markers):
continue
name_markers = [str(m).lower() for m in (match.get("name_any") or [])]
matched_role = next((m for m in name_markers if m in role_haystack), None)
if name_markers and matched_role is None:
continue
result = copy.deepcopy(playbook)
result["role"] = matched_role or ""
return result
return None


def playbook_group_id(playbook: dict[str, Any]) -> str:
"""Return the stable group id a playbook's sibling roles share."""
return str(playbook.get("id") or "")


def resolve_kernel_anchor_path(playbook: dict[str, Any]) -> str:
"""Return a stand-in ``source_file`` path for a vendor-playbook candidate.

A vendor-playbook candidate has no rewritable device source, but
downstream tooling (``kernel_optimization.py``'s CLI, in particular)
still gates on a non-empty, path-shaped ``source_file`` before it will
dispatch to a backend at all. Point that field at the task bundle's
``kernel_anchor`` file instead of leaving it empty -- resolved to an
absolute path under ``$FORGE_PATH`` when that's set (regardless of
whether the file exists on this host yet), else an absolute path under a
fixed, obviously-synthetic root so the value is still path-shaped (it
survives ``looks_like_source_path`` even when this analysis runs on a
host without the KernelForge checkout) without ever being a bare
relative string a later ``Path(...).resolve()`` could reinterpret
against an unrelated CWD.

Args:
playbook: A matched playbook entry (as returned by
``match_vendor_operator_playbook``).

Returns:
An absolute path string; never empty as long as the playbook
declares a ``kernel_anchor``, and never relative -- a relative
string here would later be reinterpreted by ``Path(...).resolve()``
against whatever the apply-stage process's CWD happens to be, not
against this bundle (PR #1191 review finding #8).
"""
anchor = str(playbook.get("kernel_anchor") or "").strip()
bundle = str(playbook.get("task_bundle") or "").strip()
if not anchor:
return ""
relative = f"{bundle}/{anchor}" if bundle else anchor
forge_root = (os.environ.get("FORGE_PATH") or "").strip()
if forge_root:
candidate = Path(forge_root) / relative
# Returned even when the file isn't there yet: this analysis host
# may lack the KernelForge checkout the apply stage will actually
# run against, but the path must still be absolute and anchored at
# a real, known root rather than silently falling through to a
# bare relative string.
return str(candidate)
# FORGE_PATH unset: still shape-check as path-like (needed to clear
# looks_like_source_path) without letting a bare relative string survive
# to be misresolved against an unrelated CWD later in the pipeline.
return str(Path("/nonexistent-forge-path") / relative)
Loading
Loading