Make a family a kind of graph, and stop paying for the DSL to decline

python/cudnn/_pygraph.py

@@ -148,6 +148,7 @@ def __init__(
         self._reserved_uids: set = set()  # user-specified uids _alloc_uid must skip
         self._ambiguous_names: set = set()  # duplicate labels: excluded from the name index
         self._candidates: Optional[List["BaseEngine"]] = None  # manifest matches + registered
+        self._facts: Dict[Any, Any] = {}  # analyzer callable -> its record; see _facts_for()
         self._backend_declined: Optional[Exception] = None  # why the backend has no entries
         self._backend_entries: Optional[List[Any]] = None  # backend_plan_entries(), once
         self._barred_names: set = set()  # deselect_engines()
@@ -182,18 +183,12 @@ def register_backend(self, engine: "BaseEngine") -> "pygraph":
         if not isinstance(eid, int):
             raise ValueError(f"engine {engine!r} must declare a stable integer engine_id (got {eid!r})")
         lo, hi = engine.owned_id_range
-        for row in MANIFEST:  # blocks are intervals, so disjointness is decidable
-            if lo < row.id_end and row.engine_id < hi:
-                # Registering the in-tree engine that OWNS the block is normal
-                # (a test or a caller pinning one specific engine); it is the
-                # same implementation the manifest would have instantiated.
-                # Anything else claiming that range would receive its plans.
-                if type(engine).__module__ == row.module and row.engine_id <= lo and hi <= row.id_end:
-                    break
+        for family in MANIFEST:  # blocks are intervals, so disjointness is decidable
+            if lo < family.id_end and family.engine_id < hi:
                 raise ValueError(
                     f"engine {engine.name!r} ({type(engine).__module__}) claims ids [{lo}, {hi}), which overlaps the "
-                    f"block [{row.engine_id}, {row.id_end}) the in-tree manifest reserves for {row.name!r} "
-                    f"({row.module}); an engine from outside the library uses ids >= {OUT_OF_TREE_ID_BASE}"
+                    f"block [{family.engine_id}, {family.id_end}) the in-tree manifest reserves for {family.name!r} "
+                    f"({family.module}); an engine from outside the library uses ids >= {OUT_OF_TREE_ID_BASE}"
                 )
         else:
             if lo < OUT_OF_TREE_ID_BASE:
@@ -919,6 +914,19 @@ def create_execution_plans(self, heuristics: Optional[List] = None) -> None:
                 "create_execution_plans() was already called on this graph; planning is one-shot — build a new graph to re-plan, or use select_plan() to switch plans"
             )
         self._backend_heuristics = heuristics  # read by backend_plan_entries()
+
+        # Finalize -> freeze -> analyze, in that order, BEFORE any ranking runs.
+        # build_operation_graph() is where the backend's own layout inference
+        # lands (a no-op for a graph that was never lowered), and it writes
+        # through object.__setattr__ precisely to bypass the freeze — so the
+        # snapshot has to come after it, not merely after validate(). Doing it
+        # here rather than inside Router keeps the epoch boundary out of
+        # overridable policy: every engine probe and the ranking then read one
+        # set of facts describing a graph that can no longer change.
+        self._finalize_backend_layout()
+        self._freeze()
+        self._attach_facts()
+
         router = self._router or default_router
         plans = list(router.plan(self, self._candidate_engines()))
         # Validate the FINAL router output: every entry must name an engine this
@@ -936,12 +944,11 @@ def create_execution_plans(self, heuristics: Optional[List] = None) -> None:
             raise cudnn_graph_not_supported(f"no engine — python or backend — proposed a plan for this graph{why}")
         self._plans = plans
         self._planning_done = True
-        self._freeze()  # plans reference the graph as-is: no mutation from here
         self._plan_index = 0
 
     def _candidate_engines(self) -> List["BaseEngine"]:
         """The python engines this graph may dispatch to: anything
-        ``register_backend()`` added, then the in-tree manifest rows whose
+        ``register_backend()`` added, then the in-tree manifest families whose
         coarse key matches. Registered engines lead because bringing your own
         engine is an explicit act; the library's own table is the default.
         (Candidate ORDER is not rank — ``heuristics_sort`` ranks.) Cached: the
@@ -959,6 +966,85 @@ def _candidate_engines(self) -> List["BaseEngine"]:
             self._candidates = list(self._backends) + from_manifest
         return self._candidates
 
+    def _finalize_backend_layout(self) -> None:
+        """Let the backend's layout inference land before the graph is frozen.
+
+        Lowering to C++ and reflecting the result back is how the GRAPH learns
+        the strides it will really execute with (the IR's own are provisional
+        row-major) — a property of the graph, not of whichever engine ends up
+        serving it. So this runs the same way whether or not the caller
+        registered an out-of-tree engine; splitting those paths is what left
+        the frozen snapshot unenforceable, since the registered-engine path
+        lowered LATER, inside the Router, and _sync_ir_shapes_from_backend
+        writes through object.__setattr__ specifically to bypass the freeze.
+
+        A failure here means "the backend cannot represent this graph", which
+        is a decline — a python engine may still serve it — recorded the way
+        backend_plan_entries() records one so the two agree.
+
+        Not routed through backend_plan_entries() itself: that additionally runs
+        the C++ plan query (~178 ms for a 1024^3 bf16 matmul), which a Router is
+        entitled never to ask for.
+        """
+        import cudnn
+
+        if not self._backend_lowerable():  # no backend lowering for this op at all
+            return
+        try:
+            self._lower_backend_graph()
+        except (cudnn.cudnnGraphNotSupportedError, RuntimeError, ImportError, AttributeError) as exc:
+            _LOG.warning("backend could not build this graph, treating as a decline: %s", exc)
+            self._backend_declined = exc
+            self._lowered_graph = None
+            self._cpp_tensors.clear()
+            self._cpp_bog_done = False
+            self._cpp_plans_created = False
+            self._backend_entries = []
+
+    def _attach_facts(self) -> None:
+        """Describe this frozen graph in each matching family's vocabulary.
+
+        Part of planning, not something a caller invokes: the graph finds its
+        own families through the manifest and hangs one record per family off
+        itself, as an optional payload. A graph no family claims carries none.
+
+        Runs AFTER _freeze() on purpose. Analyzing a graph that can still
+        change means chasing every mutation point — the layout the backend
+        infers, a dtype set between two validate() calls, a node param written
+        in place — and missing one leaves facts describing a graph that is no
+        longer there. After the freeze there is no "later" to chase.
+
+        Facts are family-scoped by construction — an SDPA fact means nothing to
+        a GEMM engine — so this is a mapping, never a union record that every
+        family would have to widen.
+        """
+        from .engines import manifest
+        from .frost.buffers import current_sm
+
+        family = manifest.family_for(self, current_sm())
+        if family is not None:
+            analyzer = manifest.resolve_analyzer(family)
+            if analyzer is not None:
+                self._facts_for(analyzer)
+
+    def _facts_for(self, analyzer):
+        """The record ``analyzer`` produced for this graph, computing it once.
+
+        Keyed by the analyzer itself, so the ranking (which resolves it from
+        ``EngineFamily.analyzer``) and the engine (which passes the callable it
+        already imports) reach ONE record with no family name to keep in sync.
+        Two extractions of the same graph is how a heuristics feature vector
+        drifts from what the kernel actually does.
+        """
+        if not self._frozen:
+            # Still mutable (a direct engine probe before planning): answer, but
+            # do not remember. Memoizing a graph that can still change is what
+            # made facts go stale across two validate() calls.
+            return analyzer(self)
+        if analyzer not in self._facts:
+            self._facts[analyzer] = analyzer(self)
+        return self._facts[analyzer]
+
     def _unlowerable_node(self) -> Optional[Node]:
         """The first node with no C++ lowering, or None when the whole graph has one.
 

python/cudnn/engines/__init__.py

@@ -25,7 +25,7 @@
     is_python_engine,
 )
 from .heuristics import heuristics_sort
-from .manifest import MANIFEST, EngineRow
+from .manifest import MANIFEST, EngineFamily
 from .router import Router, decline_types, default_router
 
 
@@ -54,7 +54,7 @@ def __getattr__(name: str):
     "decline_types",
     "heuristics_sort",
     "MANIFEST",
-    "EngineRow",
+    "EngineFamily",
     "BACKEND_ENGINE_ID_BASE",
     "CPP_OSS_ENGINE_ID_BASE",
     "PYTHON_ENGINE_ID_BASE",

python/cudnn/engines/engine_ids.py

@@ -28,12 +28,17 @@
 # learned about them in this change — so there is nothing recorded to migrate.
 PYTHON_ENGINE_ID_BASE = 20_000
 
-# Python families. A family owns a contiguous block so one family can expose
-# several ids (per kernel / per cell) without stepping on its neighbours.
-LINEAR_ATTENTION_ID_BASE = PYTHON_ENGINE_ID_BASE + 100  # 20_100..20_199
-FROST_GEMM_ID_BASE = PYTHON_ENGINE_ID_BASE + 200  # 20_200..20_299
-FROST_SDPA_FWD_ID_BASE = PYTHON_ENGINE_ID_BASE + 300  # 20_300..20_399
-FROST_SDPA_BWD_ID_BASE = PYTHON_ENGINE_ID_BASE + 400  # 20_400..20_499
+# One block per FAMILY, FAMILY_BLOCK ids wide. A family is a kind of graph (see
+# engines/manifest.py); every engine that serves that kind draws its id from
+# that family's block, taking the next free slot. A family never opens a second
+# block, so engine_id alone identifies the family and the two cannot disagree.
+FAMILY_BLOCK = 100
+GDN_ID_BASE = PYTHON_ENGINE_ID_BASE + 100  # 20_100..20_199
+KDA_ID_BASE = PYTHON_ENGINE_ID_BASE + 200  # 20_200..20_299
+GDN2_ID_BASE = PYTHON_ENGINE_ID_BASE + 300  # 20_300..20_399
+FROST_GEMM_ID_BASE = PYTHON_ENGINE_ID_BASE + 400  # 20_400..20_499
+FROST_SDPA_FWD_ID_BASE = PYTHON_ENGINE_ID_BASE + 500  # 20_500..20_599
+FROST_SDPA_BWD_ID_BASE = PYTHON_ENGINE_ID_BASE + 600  # 20_600..20_699
 OUT_OF_TREE_ID_BASE = PYTHON_ENGINE_ID_BASE + 10_000  # 30_000+, register_backend()
 
 # The delegating entry: the backend picks among candidates it holds but does not

python/cudnn/engines/heuristics.py

@@ -25,7 +25,17 @@ def heuristics_sort(graph: "pygraph", python_plans: List[PlanConfig], backend_pl
     """Rank the two sides into one list. Either side may be empty.
 
     Which python engines are in the list at all is the manifest's decision
-    (``EngineRow.matches``, where ``CUDNN_FRONTEND_ENABLE_FROST_ENGINES`` is
+    (``EngineFamily.offered``, where ``CUDNN_FRONTEND_ENABLE_FROST_ENGINES`` is
     read); a plan that reaches here has already been admitted.
+
+    Placeholder ranking, deliberately. The facts plumbing this reads FROM is
+    already in place and unused: ``graph._facts_for(analyzer)`` holds the record
+    the family's analyzer produced (planning attached it before this ran), and
+    ``manifest.family_for(graph, sm)`` names the family whose vocabulary it is
+    in. A real policy has two layers — order the family's own engines using
+    those facts, then merge that against the backend's entries on a common
+    currency (predicted time), which is the only comparison that has to work
+    across families. Neither is written yet; the seam is here so that writing
+    them does not mean re-plumbing the graph.
     """
     return python_plans + backend_plans