diff --git a/docs/python_graph_and_execution_backends.md b/docs/python_graph_and_execution_backends.md new file mode 100644 index 000000000..049ea777e --- /dev/null +++ b/docs/python_graph_and_execution_backends.md @@ -0,0 +1,130 @@ +# Python-native `cudnn.pygraph` and pluggable execution backends + +## What this is + +`cudnn.pygraph` is a Python-native graph class: graph structure (nodes, +tensors, op parameters) lives in Python with full introspection, and execution +dispatches through pluggable backends — python DSL engines and the cuDNN C++ +backend. The C++ graph builder is internal +(`cudnn._pybind_module.backend_graph`) and is reached exclusively through +lowering. + +``` +cudnn.pygraph (Python IR) → create_execution_plans() → Router → routed plan list + nodes / tensors / params (route here, PlanConfig(engine_id, knobs): + fully introspectable lazy lowering) python engines + one backend entry +``` + +Why: python-DSL engines (CuTe-DSL / cuTile style GEMM and attention fusions) +need to *see* the graph to decide whether and how to run it. Previously that +required monkey-patching the pybind class and recording calls; now the graph +is natively introspectable and an engine is one file implementing +`BaseEngine`. + +## Architecture + +### Graph IR + +- `graph_types.Tensor`, `nodes.Node`, `_pygraph.pygraph` — an engine-agnostic + op DAG. Input/output **port names equal the C++ pybind kwarg names**, + everywhere. +- Three declarative op mechanisms cover 100% of the C++ op surface: + `_POINTWISE_TENSOR_ARGS` (54 uniform pointwise ops; `mode` == method name), + `_STRUCTURED_OPS` (25 ops: norms, reduction, block-scale, MoE, conv, + structural — one table entry each: ports, attrs, outputs, shape-infer), + `_CAPTURED_OPS` (6 SDPA variants, ~130 kwargs: generic capture over an + explicit per-op schema carrying positional order, output-direction kwargs, + and conditional outputs). `matmul` is explicit for positional ergonomics. + +### Backend contract (`engines/`) + +- `BaseEngine`: `propose_plans(graph) → [PlanConfig]` (several knob configs + per engine), `build_plan(graph, plan, ctx) → CompiledPlan` (the expensive + JIT step, once per graph/plan, cached on the graph), + `CompiledPlan.execute(graph, tensor_data, ExecutionContext)` with explicit + handle/stream/workspace/overrides. Simple eager engines implement + `execute()` only. +- Every engine owns a stable `engine_id` in a reserved region + (`PYTHON_ENGINE_ID_BASE = 1 << 20`) — reproducible pinning/autotune. +- An engine declines a graph ONLY via `NotImplementedError` or + `cudnn.cudnnGraphNotSupportedError`; anything else is an engine bug and + propagates. +- `ReferenceMatmulEngine` (pure PyTorch) is the in-tree contract oracle; real + DSL engines land as separate PRs, one file each. + +### Router and the two plan-index spaces + +- The Router returns the routed plan list: python `PlanConfig` entries plus + AT MOST ONE backend delegating entry (`BACKEND_HEURISTIC_ENGINE_ID`). The + final output is validated regardless of Router implementation (registered + ids only, one sentinel max, never empty). +- **Routed space**: `graph.plans`, selected with `select_plan()`. Indices are + stable — the backend entry is one index forever and never expands in place. +- **Backend space**: the cuDNN backend's own plans, discovered per graph from + the lowered graph and addressed via the classic + `get_execution_plan_count()` / `*_plan_at_index()` APIs (pure delegation). + The frontend never statically enumerates backend engines — backend engine + sets vary by version and are discovered at plan time. +- Concrete backend engine configs as first-class routed entries need a typed + plan representation — heuristics/autotune follow-up scope, together with + ranking policy (the Router is pluggable at three levels: subclass, + per-graph `router=`, process-wide `default_router`). + +## Key invariants + +- **uid ownership**: the Python IR owns the whole uid namespace; every uid is + pushed explicitly to C++ and a post-build assertion fails loudly on + violation (C++ auto-assignment never runs for Python-built graphs — its + enumeration order is nondeterministic for multi-output ops). A user uid + landing on an auto-assigned one steals it (the holder is renumbered); + user-user collisions raise. +- **Pure-python or pure-C++**: a graph routed to a python engine never + touches C++ on the execute path; mixed construction is unsupported. + (Explicitly querying the backend plan space lowers the backend entry on + demand — that is the caller asking for the backend.) +- **One-shot planning**: a second `create_execution_plans()` raises (the + classic C++ graph never supported re-planning — it appends engine configs + by accident). Switch plans with `select_plan()`; plan differently by + building a new graph. +- **Whole-surface freeze**: after lowering/planning, every public mutation + path raises — op builders and fluent setters, direct attribute writes on + `Tensor`/`Node`/`GraphContext`, dict writes on node ports/params + (MappingProxy), in-place dim/stride edits (sealed to tuples). Inspection + stays fully readable. A mutation in the mutable window after `validate()` + invalidates the validation. +- **Output layout contract**: only USER-assigned output dim/stride are pushed + to the lowered graph; IR-inferred strides are provisional (row-major) and + the backend keeps its classic per-op layout inference (e.g. channels-last + conv). A unified layout resolver across python/cuDNN candidates belongs to + the heuristics follow-up. +- **Classic parity**: the public `cudnn.pygraph` surface behaves as before — + `cudnnGraphNotSupportedError` at `validate()`, conditional outputs return + `None`, torch dtypes/`torch.Size` accepted, ragged (THD) offsets and + multipliers on outputs, serialize/deserialize passthrough, plan queries + delegate to the lowered graph. + +## Naming + +- `cudnn.pygraph` — THE public graph class (Python IR), implemented in + `cudnn/_pygraph.py`. +- `cudnn._pybind_module.backend_graph` — the internal C++ builder the IR + lowers to (renamed from its pre-flip public name to avoid two things called + `pygraph`). + +## Testing the backend path + +The `test_native_backend_lowering.py` suite builds graphs natively, lowers, +executes on GPU, and checks numerics against torch references. Dispatch-level +assertions (`selected_engine is None`, backend plans created, lowered graph +present) prove the execution went through the cuDNN backend plan path rather +than a python engine; kernel identity below the backend API is deliberately +not asserted (kernel names are backend-internal and version-dependent). + +## Follow-ups (separate MRs) + +- Heuristics/ranking: pluggable Router policy + typed plan representation. +- DSL engine integration (the cuTile matmul engine lives in this track). +- Structural cleanup: lifecycle state objects, a `CudnnBackendAdapter` to + remove `selected_engine is None` branching, lowering extracted to its own + module, op-identity dedup (NodeType vs registry keys), longer-term a typed + `OpSpec` as the single per-op source for builder/validation/lowering. diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index ccd1b4169..43c7193ef 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -33,7 +33,6 @@ def is_windows(): "data_type", "tensor_reordering", "heur_mode", - "pygraph", "tensor", "knob", "cudnnGraphNotSupportedError", @@ -108,7 +107,7 @@ def _set_data_type( _pybind_module.tensor.set_data_type = _set_data_type -pygraph.tensor = _tensor +_pybind_module.backend_graph.tensor = _tensor def _library_device_pointer(input_tensor): @@ -194,8 +193,8 @@ def _execute_plan_at_index( ) -pygraph.execute = _execute -pygraph.execute_plan_at_index = _execute_plan_at_index +_pybind_module.backend_graph.execute = _execute +_pybind_module.backend_graph.execute_plan_at_index = _execute_plan_at_index def load_cudnn(): @@ -255,6 +254,15 @@ def _dlopen_cudnn(): else: _dlopen_cudnn() +# The graph API: a Python-native IR with pluggable execution backends. The +# public ``cudnn.pygraph`` IS the Python class; the C++ graph builder stays +# internal at ``cudnn._pybind_module.backend_graph`` and is reached only through +# lowering (a graph is pure-Python or pure-C++, never mixed). Imported before +# .graph/.wrapper, which reference cudnn.pygraph at module load. +from .graph_types import NodeType, Tensor +from ._pygraph import pygraph, GraphContext +from .nodes import Node + from .graph import graph, jit, graph_cache from .wrapper import Graph diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py new file mode 100644 index 000000000..4993ed1dd --- /dev/null +++ b/python/cudnn/_pygraph.py @@ -0,0 +1,2047 @@ +"""Pure Python graph representation for cuDNN Frontend. + +All graph structure and attributes are kept in Python. Graph construction is +backend-agnostic; a backend is chosen at create_execution_plans() time by the +Router, and the backend-specific representation (e.g. the C++ cuDNN graph) is +generated lazily only then. + +Execution flow (unification proposal): + build ops -> create_execution_plans() -> Router -> selected backend + (a registered native engine, or the cuDNN Graph backend by lazy lowering) + +Example with a native backend (pass torch tensors directly): + >>> graph = pygraph() + >>> graph.register_backend(MyDslEngine()) # any BaseEngine + >>> C = graph.matmul(a_tensor, b_tensor) # auto-creates descriptors + >>> graph.execute({C: c_tensor}) # routes to a supporting backend, else cuDNN +""" + +from dataclasses import dataclass +import weakref +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +from .graph_types import NodeType, Tensor +from .nodes import Node, _row_major_stride + +if TYPE_CHECKING: + from .engines.base import BaseEngine + + +@dataclass +class GraphContext: + """Graph-level configuration defaults.""" + + io_data_type: Any = None + intermediate_data_type: Any = None + compute_data_type: Any = None + + def __setattr__(self, name, value): + if getattr(self, "_frozen", False) and name != "_frozen": + raise RuntimeError("the graph is frozen after lowering/planning — build a new graph to change its configuration") + object.__setattr__(self, name, value) + + +class pygraph: + """Pure Python graph representation. + + All graph structure and attributes are kept in Python. C++ is only + used for execution via lazy lowering. + + Example: + >>> graph = pygraph(io_data_type=cudnn.data_type.HALF) + >>> A = graph.tensor(dim=[8, 64, 128], name="A") + >>> B = graph.tensor(dim=[8, 128, 256], name="B") + >>> C = graph.matmul(A, B, name="mm1") + >>> C.set_output(True) # outputs are explicit, like the classic API + >>> + >>> # Inspect graph + >>> print(graph.nodes) # [Node('mm1', MATMUL)] + >>> print(graph.nodes[0].inputs) # {"A": ..., "B": ...} + >>> print(graph.nodes[0].params) # {"padding": 0.0} + """ + + def __init__( + self, + # ---- classic pybind constructor, POSITIONALLY IDENTICAL (existing + # callers pass name/handle/sm_count/... by position; guarded by + # test_api_signature_parity) -------------------------------------- + name: str = "test_graph", + io_data_type: Any = None, + intermediate_data_type: Any = None, + compute_data_type: Any = None, + handle: Any = None, + sm_count: Any = None, + sm_version: Any = None, + kernel_cache: Any = None, + device_property: Any = None, + is_dynamic_shape_enabled: bool = False, + is_override_shape_enabled: bool = False, + *, + # ---- new (keyword-only: never shifts the classic positional order) -- + backends: Optional[List["BaseEngine"]] = None, + router: Any = None, + **kwargs, + ): + self._context = GraphContext( + io_data_type=io_data_type, + intermediate_data_type=intermediate_data_type or io_data_type, + compute_data_type=compute_data_type or io_data_type, + ) + self._handle = handle # cuDNN handle for the cuDNN lowering path + # Classic graph-level configuration, forwarded verbatim to the C++ + # graph at lowering (**kwargs covers future binding args). + self._cpp_graph_kwargs = {k: v for k, v in kwargs.items() if v is not None} + self._cpp_graph_kwargs["name"] = name + for _k, _v in ( + ("sm_count", sm_count), + ("sm_version", sm_version), + ("kernel_cache", kernel_cache), + ("device_property", device_property), + ): + if _v is not None: + self._cpp_graph_kwargs[_k] = _v + if is_dynamic_shape_enabled: + self._cpp_graph_kwargs["is_dynamic_shape_enabled"] = True + if is_override_shape_enabled: + self._cpp_graph_kwargs["is_override_shape_enabled"] = True + self._nodes: List[Node] = [] + self._tensors: Dict[str, Tensor] = {} + self._tensor_by_uid: Dict[int, Tensor] = {} + self._next_uid: int = 1 + self._node_count: Dict[str, int] = {} + self._lowered_graph: Any = None + self._is_validated: bool = False + self._is_built: bool = False + self._data_bindings: Dict[int, Any] = {} # uid -> tensor data for auto-bound inputs + + # Backend routing (see engines/router.py). Graph construction is + # backend-agnostic. At create_execution_plans() the Router builds a flat + # ranked plan list (python engines + cuDNN) in one shared engine-id + # space; each plan is dispatched by its id (is_python_engine -> python + # registry, else lower to cuDNN). ``_plan_index`` selects the plan to run. + self._backends: List["BaseEngine"] = [] + self._router = router # None => engines.router.default_router at route time + self._plans: List[Any] = [] # list[PlanConfig], populated by create_execution_plans() + self._planning_done: bool = False # create_execution_plans() ran (one-shot) + self._frozen: bool = False # whole-surface freeze (set by _freeze()) + self._plan_index: int = 0 + self._backend_heuristics: Optional[List] = None # heur modes for a backend plan + self._cpp_plans_created: bool = False # C++ create_execution_plans ran + self._compiled_plans: Dict[int, Any] = {} # plan_index -> CompiledPlan (python plans) + self._cpp_bog_done: bool = False # C++ build_operation_graph ran + self._cpp_tensors: Dict[int, Any] = {} # IR uid -> lowered C++ tensor + self._reserved_uids: set = set() # user-specified uids _alloc_uid must skip + self._ambiguous_names: set = set() # duplicate labels: excluded from the name index + for _e in backends or (): # constructor path uses the SAME validation + self.register_backend(_e) + + # ========================================================================= + # Backend registration & routing + # ========================================================================= + + def register_backend(self, engine: "BaseEngine") -> "pygraph": + """Add a candidate python execution engine. It joins the plan list at + create_execution_plans() time when its check_support() accepts the graph. + + Validated at registration (not at failure time): the engine must declare + a stable engine_id in the reserved python region, ids must be unique per + graph, and registration after planning is rejected (planning is + one-shot — build a new graph).""" + from .engines.engine_ids import is_python_engine + + eid = getattr(engine, "engine_id", None) + if not isinstance(eid, int) or not is_python_engine(eid): + raise ValueError(f"engine {engine!r} must declare a stable integer engine_id >= PYTHON_ENGINE_ID_BASE (got {eid!r})") + if any(e.engine_id == eid for e in self._backends): + raise ValueError(f"engine_id {eid} is already registered on this graph") + if self._planning_done: + raise RuntimeError("cannot register a backend after create_execution_plans(); planning is one-shot — build a new graph") + self._backends.append(engine) + return self + + def set_router(self, router: Any) -> "pygraph": + """Override the plan-list / ranking policy for this graph. Must be set + before create_execution_plans() (a later router cannot affect the + already-planned list).""" + if self._planning_done: + raise RuntimeError("cannot set a router after create_execution_plans(); planning is one-shot — build a new graph") + self._router = router + return self + + def _engine_by_id(self, engine_id: int) -> "BaseEngine": + for e in self._backends: + if e.engine_id == engine_id: + return e + raise KeyError(f"no registered python engine with id {engine_id}") + + @property + def backends(self) -> List["BaseEngine"]: + """Registered candidate python engines.""" + return list(self._backends) + + @property + def plans(self) -> List[Any]: + """The ranked plan list (list[PlanConfig]) from create_execution_plans().""" + return list(self._plans) + + @property + def _selected_plan_config(self) -> Optional[Any]: + from .engines.engine_ids import is_python_engine + + if not self._plans or not 0 <= self._plan_index < len(self._plans): + return None + cfg = self._plans[self._plan_index] + return cfg if is_python_engine(cfg.engine_id) else None + + @property + def selected_engine(self) -> Optional["BaseEngine"]: + """The python engine for the currently selected top-level plan entry, + or None for the backend path. Populated after create_execution_plans().""" + cfg = self._selected_plan_config + return self._engine_by_id(cfg.engine_id) if cfg is not None else None + + # ========================================================================= + # Tensor Creation + # ========================================================================= + + def tensor( + self, + # classic public tensor() signature, POSITIONALLY IDENTICAL (guarded by + # test_api_signature_parity); classic unset sentinels (NOT_SET, -1, + # reordering NONE) are normalized to None below + dim: List[int], + stride: Optional[List[int]] = None, + data_type: Any = None, + is_virtual: bool = False, + is_pass_by_value: bool = False, + ragged_offset: Optional[Tensor] = None, + reordering_type: Any = None, + name: str = "", + uid: Optional[int] = None, + ragged_offset_multiplier: int = 1, + **kwargs, + ) -> Tensor: + """Create a tensor.""" + if not name: + name = f"tensor_{len(self._tensors)}" + if data_type is not None and getattr(data_type, "name", None) == "NOT_SET": + data_type = None + if reordering_type is not None and getattr(reordering_type, "name", None) == "NONE": + reordering_type = None + if uid == -1: # classic unset sentinel + uid = None + + if uid is not None: + # User-owned uid, same rule as set_uid (_reuid_tensor): classic + # tensors have no uid until assigned, so a user uid may land on an + # eagerly auto-assigned one — the user wins, the auto holder is + # renumbered; colliding with another USER uid is an error. + holder = self._tensor_by_uid.get(uid) + if holder is not None: + if holder.uid_assigned: + raise ValueError(f"uid {uid} is already user-assigned to tensor {holder.name!r}") + fresh = self._alloc_uid() + self._tensor_by_uid[fresh] = holder + if holder.uid in self._data_bindings: + self._data_bindings[fresh] = self._data_bindings.pop(holder.uid) + holder.uid = fresh + self._reserved_uids.add(uid) + + dim = list(dim) # classic API accepts torch.Size / tuples + t = Tensor( + name=name, + dim=dim, + stride=list(stride) if stride else _row_major_stride(dim), + data_type=data_type or (self._context.intermediate_data_type if is_virtual else self._context.io_data_type), + is_virtual=is_virtual, + is_pass_by_value=is_pass_by_value, + ragged_offset=ragged_offset, + reordering_type=reordering_type, + ragged_offset_multiplier=ragged_offset_multiplier, + uid=uid if uid is not None else self._alloc_uid(), + uid_assigned=uid is not None, + dim_assigned=True, # graph inputs: the user specified the layout + stride_assigned=stride is not None, + **kwargs, + ) + self._register_tensor(t) + return t + + def tensor_like(self, template: Any, name: str = "", is_virtual: bool = False) -> Tensor: + """Create tensor from another IR tensor or a DLPack object (e.g. torch). + + Classic parity: CPU (host) framework tensors become pass-by-value, like + the C++ tensor_like (is_pass_by_value = device == CPU).""" + if isinstance(template, Tensor): + return self.tensor( + dim=list(template.dim), + stride=list(template.stride), + data_type=template.data_type, + is_virtual=is_virtual, + is_pass_by_value=template.is_pass_by_value, + name=name, + ) + # Element strides via the DLPack protocol (classic tensor_like reads the + # DLPack capsule). torch's .stride() is element units, but e.g. CuPy + # exposes byte-unit .strides — so normalize any non-torch DLPack object + # through torch.from_dlpack first. + if hasattr(template, "stride") and callable(getattr(template, "stride", None)): + dim = list(template.shape) + stride = list(template.stride()) + else: + try: + import torch as _torch + + _view = _torch.from_dlpack(template) + dim = list(_view.shape) + stride = list(_view.stride()) + except Exception: # noqa: BLE001 — no torch / exotic dlpack: assume dense + dim = list(template.shape) + stride = _row_major_stride(dim) + + data_type = None + try: + import cudnn.datatypes + + data_type = cudnn.datatypes._torch_to_cudnn_data_type(template.dtype) + except Exception: + pass + + is_pbv = bool(getattr(getattr(template, "device", None), "type", None) == "cpu") + return self.tensor(dim=dim, stride=stride, data_type=data_type, is_virtual=is_virtual, is_pass_by_value=is_pbv, name=name) + + def tensor_scalar(self, value: Any, scalar_type: Any, name: str = "") -> Tensor: + """Create a pass-by-value scalar tensor (classic tensor_scalar parity).""" + if not name: + name = f"scalar_{len(self._tensors)}" + t = Tensor( + name=name, + dim=[1, 1, 1, 1], + stride=[1, 1, 1, 1], + is_pass_by_value=True, + pass_by_value=value, + scalar_type=scalar_type, + uid=self._alloc_uid(), + ) + self._register_tensor(t) + return t + + def _check_mutable(self, what: str) -> None: + if self._frozen: + raise RuntimeError(f"cannot {what} after lowering/planning — the graph is frozen (planning is one-shot; build a new graph)") + # a mutation while merely validated (python-engine graphs stay mutable + # until planning) must re-validate later — never run on stale inference + self._is_validated = False + + def _freeze(self) -> None: + """Freeze the ENTIRE public graph surface (not just the fluent API). + + Called at lowering and at planning, whichever happens first. After + this, every mutation path raises: fluent setters and op builders (via + _check_mutable), attribute writes on Tensor/Node/GraphContext (their + __setattr__ guards), dict writes on node.inputs/outputs/params + (MappingProxy), and in-place list mutation of dim/stride (tuples). + The inspection surface stays fully readable for engines.""" + if self._frozen: + return + from types import MappingProxyType + + for node in self._nodes: + node.inputs = MappingProxyType(dict(node.inputs)) + node.outputs = MappingProxyType(dict(node.outputs)) + node.params = MappingProxyType(dict(node.params)) + node._frozen = True + for t in self._tensor_by_uid.values(): + t.dim = tuple(t.dim) if t.dim else t.dim + t.stride = tuple(t.stride) if t.stride else t.stride + t._frozen = True + self._context._frozen = True + self._frozen = True + + def _rename_tensor(self, t: Tensor, name: str) -> None: + """Atomic rename keeping the name index coherent. Classic parity: names + are labels, so renaming ONTO an existing label is legal — the name just + becomes ambiguous and leaves the unique-name index.""" + if name == t.name: + return + # NOT freeze-guarded: names are labels (classic allows renaming after + # build — the lowered graph already carries the old label, and labels + # have no execution semantics). + if self._tensors.get(t.name) is t: + del self._tensors[t.name] + object.__setattr__(t, "name", name) # label write is exempt from the freeze + if name in self._tensors or name in self._ambiguous_names: + self._tensors.pop(name, None) + self._ambiguous_names.add(name) + else: + self._tensors[name] = t + + def _reuid_tensor(self, t: Tensor, uid: int) -> None: + """Atomic re-uid keeping indexes/bindings coherent. + + Classic parity: classic tensors have NO uid until set_uid, while the IR + assigns eagerly — so a user set_uid may land on an auto-assigned uid. + The user wins: the auto holder is silently renumbered (auto uids are + internal until lowering). Two USER-assigned uids colliding is an error. + """ + if uid == t.uid: + if not self._frozen: # same-value set_uid is a no-op (classic allows it anytime) + t.uid_assigned = True + return + self._check_mutable("re-uid a tensor") + holder = self._tensor_by_uid.get(uid) + if holder is not None: + if holder.uid_assigned: + raise ValueError(f"uid {uid} is already user-assigned to tensor {holder.name!r}") + fresh = self._alloc_uid() # renumber the auto holder + self._tensor_by_uid[fresh] = holder + if holder.uid in self._data_bindings: + self._data_bindings[fresh] = self._data_bindings.pop(holder.uid) + holder.uid = fresh + self._tensor_by_uid.pop(t.uid, None) + if t.uid in self._data_bindings: + self._data_bindings[uid] = self._data_bindings.pop(t.uid) + t.uid = uid + t.uid_assigned = True + self._reserved_uids.add(uid) + self._tensor_by_uid[uid] = t + + def _alloc_uid(self) -> int: + # Skip uids the user reserved via tensor(uid=...) — the Python IR owns + # the whole uid namespace (see the uid-ownership note in _lower_to_cpp). + while self._next_uid in self._reserved_uids: + self._next_uid += 1 + uid = self._next_uid + self._next_uid += 1 + return uid + + # Classic C++ auto-names op outputs "::" (graph_interface.h + # output_tensor calls). Ports whose name differs from the classic enum are + # mapped here so canonical names (wrapper.Graph lookups, JSON dumps) match. + _CLASSIC_OUT_SUFFIX = { + "inv_var": "INV_VARIANCE", + "mean": "MEAN", + "next_running_mean": "NEXT_RUNNING_MEAN", + "next_running_var": "NEXT_RUNNING_VAR", + "DScale": "DSCALE", + "DBias": "DBIAS", + } + + def _get_name(self, op: str, name: str) -> str: + self._check_mutable(f"add a {op} op") + if name: + return name + count = self._node_count.get(op, 0) + self._node_count[op] = count + 1 + return f"{op}.{count}" + + def _make_output(self, name: str) -> Tensor: + """Create a virtual output tensor. data_type is left unset: validate()'s + inference assigns io/intermediate by the FINAL virtual state (classic + semantics — a user set_output(True) without set_data_type gets io).""" + return Tensor( + name=name, + is_virtual=True, + uid=self._alloc_uid(), + ) + + def _register_tensor(self, t: Tensor) -> None: + t.owner = weakref.ref(self) + # Classic parity: names are debug LABELS — duplicates are legal + # (pycudnnTest builds two 'weight' tensors). uid is the identity; the + # name index serves only names that remain unique, and name-keyed + # lookups on an ambiguous name raise instead of guessing. + if t.name in self._tensors or t.name in self._ambiguous_names: + self._tensors.pop(t.name, None) + self._ambiguous_names.add(t.name) + else: + self._tensors[t.name] = t + self._tensor_by_uid[t.uid] = t + + def _ensure_tensor(self, arg: Any, name: str = "") -> Tensor: + """Convert arg to a Tensor descriptor if it isn't one already. + + If arg is a framework tensor (torch, jax, cupy, etc.), creates a + descriptor via tensor_like() and stores the data binding for execute(). + """ + if isinstance(arg, Tensor): + return arg + desc = self.tensor_like(arg, name=name) + self._data_bindings[desc.uid] = arg + return desc + + # ========================================================================= + # Operations + # ========================================================================= + + def matmul( + self, + A: Any, + B: Any, + compute_data_type: Any = None, + padding: float = 0.0, + name: str = "", + ) -> Tensor: + """Matrix multiplication: C = A @ B. + + A and B can be Tensor descriptors or framework tensors (torch, jax, etc.). + """ + name = self._get_name("matmul", name) + A = self._ensure_tensor(A, name=f"{name}::A") + B = self._ensure_tensor(B, name=f"{name}::B") + + node = Node(name, NodeType.MATMUL, compute_data_type or self._context.compute_data_type) + node.inputs["A"] = A + node.inputs["B"] = B + node.params["padding"] = padding + + C = self._make_output(f"{name}::C") + node.outputs["C"] = C + self._register_tensor(C) + + self._nodes.append(node) + return C + + # ---- Pointwise ops ------------------------------------------------------ + # ``params["mode"]`` is the op kind == the C++ pygraph method name (the + # pointwise_mode enum is not exposed to Python, and the method name IS the + # canonical semantic name), so lowering is a direct getattr dispatch — no + # mode<->method mapping table to maintain. Extra scalar attributes + # (negative_slope / clips / swish_beta / axis) live in params and are + # forwarded at lowering; ops that take them get explicit builders below, + # the uniform rest are generated from _POINTWISE_TENSOR_ARGS (the table + # mirrors the pybind signatures — tensor-argument names per op — so both + # positional and the classic keyword call styles work). + + _POINTWISE_TENSOR_ARGS: "dict[str, tuple]" = { + # unary + **{ + op: ("input",) + for op in ( + "abs", + "ceil", + "cos", + "elu", + "erf", + "exp", + "floor", + "gelu", + "gelu_approx_tanh", + "identity", + "log", + "logical_not", + "neg", + "reciprocal", + "rsqrt", + "sigmoid", + "sin", + "softplus", + "sqrt", + "tan", + "tanh", + ) + }, + # binary + **{op: ("a", "b") for op in ("add", "add_square", "div", "logical_and", "logical_or", "mul", "sub")}, + **{op: ("input0", "input1") for op in ("max", "min", "mod", "pow")}, + **{op: ("input", "comparison") for op in ("cmp_eq", "cmp_ge", "cmp_gt", "cmp_le", "cmp_lt", "cmp_neq")}, + "bias": ("input", "bias"), + "scale": ("input", "scale"), + # backward (loss, input) -> dinput + **{ + op: ("loss", "input") + for op in ( + "elu_backward", + "gelu_approx_tanh_backward", + "gelu_backward", + "sigmoid_backward", + "softplus_backward", + "tanh_backward", + ) + }, + # ternary + "binary_select": ("input0", "input1", "mask"), + } + # scalar attributes forwarded from params to the C++ call at lowering + _POINTWISE_EXTRA_PARAMS = ("negative_slope", "lower_clip", "upper_clip", "swish_beta", "axis") + + def _pointwise(self, mode: str, inputs: list, name: str, compute_data_type: Any = None, extra_params: Optional[dict] = None) -> Tensor: + """Internal helper for pointwise ops. ``mode`` == C++ pygraph method name.""" + inputs = [self._ensure_tensor(t, name=f"{name}::IN_{i}") for i, t in enumerate(inputs)] + node = Node(name, NodeType.POINTWISE, compute_data_type or self._context.compute_data_type) + node.params["mode"] = mode + if extra_params: + node.params.update({k: v for k, v in extra_params.items() if v is not None}) + for i, t in enumerate(inputs): + node.inputs[f"IN_{i}"] = t + + out = self._make_output(f"{name}::OUT_0") + node.outputs["OUT_0"] = out + self._register_tensor(out) + + self._nodes.append(node) + return out + + # Pointwise ops with extra scalar attributes: explicit builders. + + def relu( + self, input: Any, negative_slope: Any = None, lower_clip: Any = None, upper_clip: Any = None, name: str = "", compute_data_type: Any = None + ) -> Tensor: + """ReLU (optionally leaky via negative_slope, and/or clipped).""" + return self._pointwise( + "relu", [input], self._get_name("relu", name), compute_data_type, dict(negative_slope=negative_slope, lower_clip=lower_clip, upper_clip=upper_clip) + ) + + def leaky_relu(self, input: Any, negative_slope: Any, name: str = "", compute_data_type: Any = None) -> Tensor: + """Leaky ReLU.""" + return self._pointwise("leaky_relu", [input], self._get_name("leaky_relu", name), compute_data_type, dict(negative_slope=negative_slope)) + + def swish(self, input: Any, swish_beta: Any = None, name: str = "", compute_data_type: Any = None) -> Tensor: + """Swish / SiLU.""" + return self._pointwise("swish", [input], self._get_name("swish", name), compute_data_type, dict(swish_beta=swish_beta)) + + def gen_index(self, input: Any, axis: int, name: str = "", compute_data_type: Any = None) -> Tensor: + """Generate index along an axis.""" + return self._pointwise("gen_index", [input], self._get_name("gen_index", name), compute_data_type, dict(axis=axis)) + + def relu_backward( + self, loss: Any, input: Any, negative_slope: Any = None, lower_clip: Any = None, upper_clip: Any = None, name: str = "", compute_data_type: Any = None + ) -> Tensor: + """ReLU backward.""" + return self._pointwise( + "relu_backward", + [loss, input], + self._get_name("relu_backward", name), + compute_data_type, + dict(negative_slope=negative_slope, lower_clip=lower_clip, upper_clip=upper_clip), + ) + + def leaky_relu_backward(self, loss: Any, input: Any, negative_slope: Any, name: str = "", compute_data_type: Any = None) -> Tensor: + """Leaky ReLU backward.""" + return self._pointwise( + "leaky_relu_backward", [loss, input], self._get_name("leaky_relu_backward", name), compute_data_type, dict(negative_slope=negative_slope) + ) + + def swish_backward(self, loss: Any, input: Any, swish_beta: Any = None, name: str = "", compute_data_type: Any = None) -> Tensor: + """Swish backward.""" + return self._pointwise("swish_backward", [loss, input], self._get_name("swish_backward", name), compute_data_type, dict(swish_beta=swish_beta)) + + # NOTE: reduction / block-scale / MoE / conv / norms / structural ops are all + # declared in _STRUCTURED_OPS (module tail) — one table entry per op, one + # generic lowering branch. Only ops whose call shape doesn't fit the table + # (matmul's positional ergonomics, sdpa's conditional kwargs) stay explicit. + + # NOTE: the sdpa family (sdpa / sdpa_backward / sdpa_fp8 / sdpa_mxfp8 / + # sdpa_fp8_backward / sdpa_mxfp8_backward) is declared in _CAPTURED_OPS + # (module tail): kwargs are captured generically — tensors become named + # ports (port == C++ kwarg), scalars/enums/callbacks go to params verbatim, + # dropout tuples are flattened per element — and lowering forwards them + # verbatim, so the full C++ kwarg surface (~130 args across variants) is + # supported without hand-mirroring each argument. + + # ========================================================================= + # Inspection + # ========================================================================= + + @property + def nodes(self) -> List[Node]: + """All nodes in the graph (a copy — the graph's own list is not a + public mutation path).""" + return list(self._nodes) + + @property + def tensors(self) -> Dict[str, Tensor]: + """All tensors by name (a copy — see nodes).""" + return dict(self._tensors) + + @property + def context(self) -> GraphContext: + """Graph context.""" + return self._context + + def find_tensor(self, name_or_uid: Union[str, int]) -> Optional[Tensor]: + """Find tensor by name or UID.""" + if isinstance(name_or_uid, int): + return self._tensor_by_uid.get(name_or_uid) + if name_or_uid in self._ambiguous_names: + raise ValueError(f"tensor name {name_or_uid!r} is ambiguous (duplicate labels are legal; look up by uid or Tensor)") + return self._tensors.get(name_or_uid) + + def get_node(self, name: str) -> Optional[Node]: + """Find node by name.""" + return next((n for n in self._nodes if n.name == name), None) + + def get_inputs(self) -> List[Tensor]: + """Get non-virtual input tensors.""" + produced = {t.uid for n in self._nodes for t in n.outputs.values() if t} + return [t for t in self._tensors.values() if not t.is_virtual and t.uid not in produced] + + def get_outputs(self) -> List[Tensor]: + """Get non-virtual output tensors.""" + return [t for t in self._tensors.values() if not t.is_virtual and any(t.uid == o.uid for n in self._nodes for o in n.outputs.values() if o)] + + def inspect(self) -> Dict[str, Any]: + """Return graph structure for inspection.""" + return { + "context": { + "io_data_type": str(self._context.io_data_type), + "compute_data_type": str(self._context.compute_data_type), + }, + "nodes": [ + { + "name": n.name, + "type": n.node_type.name, + "inputs": {k: v.name for k, v in n.inputs.items()}, + "outputs": {k: v.name for k, v in n.outputs.items()}, + "params": n.params, + } + for n in self._nodes + ], + "tensors": { + name: {"dim": t.dim, "stride": t.stride, "dtype": str(t.data_type), "is_virtual": t.is_virtual, "uid": t.uid} + for name, t in self._tensors.items() + }, + } + + # ========================================================================= + # Build & Execute + # ========================================================================= + + def validate(self) -> None: + """Validate graph and infer properties. + + Classic parity: op outputs stay VIRTUAL unless the user marks them + with set_output(True). A leaf output is NOT auto-marked — discarding + an op result (e.g. the Stats of a training SDPA) is legal classic + usage, and auto-marking it would make its uid required in the variant + pack. + """ + for node in self._nodes: + node.infer_properties(self._context) + # Table-driven shape inference, topologically: builder-time infer + # only sees graph-input dims; chained ops (e.g. conv on a virtual + # relu output) get their output dims here, once inputs are known. + spec_entry = _STRUCTURED_BY_TYPE.get(node.node_type) or _CAPTURED_BY_TYPE.get(node.node_type) + if spec_entry: + _, spec = spec_entry + infer = spec.get("infer", {}) + for oport, out_t in node.outputs.items(): + if out_t is not None and not out_t.dim: + try: + d = infer.get(oport, lambda n: None)(node) + except Exception: # noqa: BLE001 — best-effort + d = None + if d: + out_t.dim = list(d) + out_t.stride = _row_major_stride(out_t.dim) + node.validate() + for t in self._tensors.values(): + if t.dim and not t.stride: # classic: stride optional, row-major inferred + t.stride = _row_major_stride(t.dim) + if not t.is_pass_by_value: + t.validate() + self._is_validated = True + # Classic parity: with no python engines registered, C++ validation + # happens HERE — tests catch cudnnGraphNotSupportedError around + # graph.validate() (unsupported configs must skip, not fail later). + if not self._backends and self._lowered_graph is None: + self._lowered_graph = self._lower_to_cpp() + self._lowered_graph.validate() + self._verify_uid_ownership() + + def build_operation_graph(self) -> None: + """Validate the graph; lower to C++ when no python engines are registered. + + Backend selection is deferred to create_execution_plans() (the Router + stage). With python engines registered, nothing is lowered here (a graph + routed to a python engine never touches C++). Without them — the classic + sequencing — lowering happens now, so plan-configuration and query + methods (deselect_engines, get_engine_count, ...) work between + build_operation_graph() and create_execution_plans(), exactly as on the + classic API (they delegate to the lowered C++ graph via __getattr__). + """ + if not self._is_validated: + self.validate() + if self._lowered_graph is not None and not self._cpp_bog_done: + self._lowered_graph.build_operation_graph() + self._cpp_bog_done = True + self._sync_ir_shapes_from_backend() + + def _sync_ir_shapes_from_backend(self) -> None: + """After the backend's shape/layout inference (build_operation_graph), + reflect the REAL dim/stride back into the IR tensors. The IR's own + inferred strides are provisional row-major; the backend applies + classic per-op layout inference (channels-last conv etc.), and + consumers of the IR (wrapper.Graph buffer allocation, engines, + introspection) must see the layout that will actually execute.""" + for ir_uid, cpp_t in self._cpp_tensors.items(): + ir = self._tensor_by_uid.get(ir_uid) + if ir is None: + continue + try: + d, st = cpp_t.get_dim(), cpp_t.get_stride() + except Exception: # noqa: BLE001 — some tensors have no dims (scalars) + continue + if d: + object.__setattr__(ir, "dim", tuple(d)) # sealed (graph is frozen) + if st: + object.__setattr__(ir, "stride", tuple(st)) + + def create_execution_plans(self, heuristics: Optional[List] = None) -> None: + """Build the ranked execution-plan list (the dispatch stage). + + The Router returns one flat list of PlanConfig(engine_id, knobs) mixing + python engines (reserved id region) and the backend side, in one shared + engine-id space. Nothing is lowered here — a plan is built lazily when + selected. ``_plan_index`` selects which plan runs (default 0, the + highest-ranked); cuDNN heuristic modes are carried on the backend plan's + knobs. + + Args: + heuristics: cuDNN heuristic modes, carried to the backend plan. + """ + if not self._is_validated: + self.validate() + + from .engines.router import default_router + + # One-shot planning (classic conformance: the C++ graph never supported + # re-planning — a second call there appends plans by accident, and no + # user re-plans). Plan once; to plan differently, build a new graph + # (IR construction is microseconds). Autotune re-selects WITHIN this + # plan set via select_plan(). Explicit state flag, not an + # is-the-list-nonempty proxy. + if self._planning_done: + raise RuntimeError( + "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" + ) + router = self._router or default_router + plans = router.plan(self, self._backends) + # Validate the FINAL router output (a custom Router must not bypass + # registration): python entries must name registered engines; the only + # non-python entry allowed is ONE backend delegating sentinel. + from .engines.engine_ids import BACKEND_HEURISTIC_ENGINE_ID, is_python_engine + + registered = {e.engine_id for e in self._backends} + if not plans: + raise ValueError("router returned an empty plan list — there is no legal empty planning state (return the backend delegating entry at minimum)") + n_cudnn = 0 + for cfg in plans: + if is_python_engine(cfg.engine_id): + if cfg.engine_id not in registered: + raise ValueError(f"router produced a plan for unregistered engine_id {cfg.engine_id}") + elif cfg.engine_id == BACKEND_HEURISTIC_ENGINE_ID: + n_cudnn += 1 + else: + raise ValueError(f"router produced a plan with invalid engine_id {cfg.engine_id}") + if n_cudnn > 1: + raise ValueError("router produced more than one backend delegating entry") + self._plans = plans + self._planning_done = True + self._freeze() # plans reference the graph as-is: no mutation from here + self._plan_index = 0 + self._backend_heuristics = heuristics # applied when a backend plan is built + # Classic sequencing: if the graph was already lowered (no python + # engines -> build_operation_graph lowered eagerly) and the selected + # plan is the backend one, create the C++ plans now. + if self.selected_engine is None and self._lowered_graph is not None: + self._lower_backend_plan() + + def _has_backend_plan(self) -> bool: + from .engines.engine_ids import BACKEND_HEURISTIC_ENGINE_ID + + return any(cfg.engine_id == BACKEND_HEURISTIC_ENGINE_ID for cfg in self._plans) + + def get_execution_plan_count(self) -> int: + """Classic passthrough, ALWAYS: the cuDNN backend's plan count for this + graph (its plan list is discovered per graph from the lowered C++ graph + and addressed via the classic ``build_plan_at_index`` / + ``execute_plan_at_index`` / ``get_workspace_size_plan_at_index`` APIs). + The semantics never depend on whether python engines are registered. + + The ROUTED plan list (the Router's entries: python plans + at most one + backend delegating entry) is a separate index space: ``graph.plans``, + selected with ``select_plan()``. Its indices are stable — the cuDNN + entry is one index forever and never expands into this count. + """ + if self._planning_done: + if not self._has_backend_plan(): + raise RuntimeError( + "this graph's Router produced python plans only (no backend entry), so there are no backend plans — the routed plan list is graph.plans / select_plan()" + ) + self._lower_backend_plan() # backend plans exist on demand (one-shot) + return self._lowered_graph.get_execution_plan_count() + if self._lowered_graph is not None: + # classic pre-planning sequencing: delegate, C++ reports its state + return self._lowered_graph.get_execution_plan_count() + return 0 # classic: an unplanned graph has zero plans (not an error) + + def select_plan(self, index: int) -> "pygraph": + """Pick a ROUTED plan entry: the index is into ``graph.plans`` (the + Router's entries — stable, never shifted by backend lowering). Backend + sub-plans are a separate space, selected via the classic at-index APIs + (see get_execution_plan_count).""" + if not self._planning_done: + raise RuntimeError("call create_execution_plans() before select_plan()") + if not 0 <= index < len(self._plans): + raise IndexError(f"plan index {index} out of range for {len(self._plans)} routed plan(s) (graph.plans)") + self._plan_index = index + self._is_built = False + return self + + def _resolve_stream(self, handle: Any) -> Any: + """Stream for a supplied handle (classic set_stream semantics). A failed + query on a SUPPLIED handle is a correctness error and raises — never a + silent fall-back to another stream. No handle -> None (the engine must + resolve deterministically from its framework, e.g. torch current stream).""" + if handle is None: + return None + import cudnn + + return cudnn.get_stream(handle) + + def _build_context(self, handle: Any = None) -> Any: + from .engines.base import ExecutionContext + + h = handle if handle is not None else self._handle + return ExecutionContext(handle=h, stream=self._resolve_stream(h)) + + def _verify_uid_ownership(self) -> None: + # Verify the uid-ownership invariant (see _lower_to_cpp): every C++ + # tensor must carry exactly its IR uid. An assertion — not a silent + # translation — so a lowering path that forgets to push a uid fails + # loudly in tests instead of mis-binding buffers (a swapped + # multi-output pairing writes past the smaller buffer: corruption). + for ir_uid, cpp_t in self._cpp_tensors.items(): + cpp_uid = cpp_t.get_uid() + if cpp_uid != ir_uid: + raise RuntimeError(f"uid ownership violated: IR tensor uid {ir_uid} lowered to C++ uid {cpp_uid} — a lowering path failed to push the uid") + + def _lower_backend_plan(self) -> None: + """Lower to C++ (if not already) and create the backend plans (once).""" + import cudnn + + if self._lowered_graph is None: + self._lowered_graph = self._lower_to_cpp() + self._lowered_graph.validate() + self._verify_uid_ownership() + if not self._cpp_bog_done: + self._lowered_graph.build_operation_graph() + self._cpp_bog_done = True + self._sync_ir_shapes_from_backend() + if not self._cpp_plans_created: + heur = self._backend_heuristics or [cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK] + self._lowered_graph.create_execution_plans(heur) + self._cpp_plans_created = True + + def check_support(self) -> None: + """Check the selected plan's engine supports the graph. + + A python plan re-affirms its engine's check_support() (already passed + when the Router included it); a backend plan lowers and checks C++ support. + """ + eng = self.selected_engine + if eng is not None: + eng.check_support(self) + return + if self._lowered_graph is None: + self._lower_backend_plan() + self._lowered_graph.check_support() + + def build_plans(self, *args) -> None: + """Finalize the selected plan. A python plan compiles HERE (once per + graph/plan; the CompiledPlan is cached on the graph and reused across + executions). The classic optional build_plan_policy passes through on + the backend path.""" + eng = self.selected_engine + if eng is not None: + if self._plan_index not in self._compiled_plans: + self._compiled_plans[self._plan_index] = eng.build_plan(self, self._selected_plan_config, self._build_context()) + if eng is None: + if self._lowered_graph is None or not self._cpp_plans_created: + self._lower_backend_plan() + self._lowered_graph.build_plans(*args) + self._is_built = True + + def build(self, heuristics: Optional[List] = None) -> None: + """Convenience: validate -> build_operation_graph -> create_execution_plans + -> check_support -> build_plans, in sequence.""" + if not self._is_validated: + self.validate() + + self.build_operation_graph() + if not self._planning_done: # never silently re-plan: preserves select_plan() + self.create_execution_plans(heuristics) + self.check_support() + self.build_plans() + + def get_workspace_size(self, *args, **kwargs) -> int: + """Workspace bytes for the selected plan. Classic overloads (handle / + dynamic-shape overrides) pass through on the backend path.""" + if not self._is_built: + raise RuntimeError("Call build() first") + + if self.selected_engine is not None: + if args or kwargs: + raise NotImplementedError("dynamic workspace-query overrides are not supported by python plans") + return self._compiled_plans[self._plan_index].get_workspace_size() + + return self._lowered_graph.get_workspace_size(*args, **kwargs) + + def execute( + self, + tensor_dict: Dict[Union[str, int, Tensor], Any], + workspace: Any = None, + handle: int = None, + override_uids: Any = None, + override_shapes: Any = None, + override_strides: Any = None, + ) -> None: + """Execute the selected plan. + + Both python engines and the backend path write results directly into the + caller-provided output tensors (in-place). Automatically calls build() + if it hasn't run yet. Dispatch is a single check on the plan's engine id. + + Args: + tensor_dict: Dict mapping tensors (by Tensor, name, or uid) to data. + Must include both input and output tensors. + workspace: Workspace buffer (ignored by python engines) + handle: cuDNN handle (ignored by python engines) + override_uids/shapes/strides: dynamic-shape overrides (backend path) + """ + if not self._is_built: + # Auto-build. When a python plan will run and the caller supplied a + # handle HERE, the JIT compile must see it — plan first, then let + # the python branch below compile with the caller's context instead + # of running the generic build (which only knows the graph handle). + if not self._planning_done: + self.create_execution_plans() + if self.selected_engine is None: + self.build() + + # Start with auto-bound inputs, then overlay user-provided (user wins) + uid_to_data = dict(self._data_bindings) + for key, data in tensor_dict.items(): + if key is None: + continue # classic API tolerates None keys (optional tensors) + if isinstance(key, Tensor): + uid = key.uid + elif isinstance(key, str): + if key in self._ambiguous_names: + raise ValueError(f"tensor name {key!r} is ambiguous (duplicate labels); key the variant pack by uid or Tensor") + uid = self._tensors[key].uid + elif isinstance(key, int): + uid = key + else: # a lowered C++ tensor (advanced/interop) — trust its uid + uid = key.get_uid() + uid_to_data[uid] = data + + eng = self.selected_engine + if eng is not None: # python engine (plan id in the reserved region) + from .engines.base import ExecutionContext + + h = handle if handle is not None else self._handle + ctx = ExecutionContext( + handle=h, + stream=self._resolve_stream(h), + workspace=workspace, + override_uids=override_uids, + override_shapes=override_shapes, + override_strides=override_strides, + ) + if self._plan_index not in self._compiled_plans: + # compile with the CALLER's context (execute-supplied handle + # and its stream reach the JIT build) + self._compiled_plans[self._plan_index] = eng.build_plan(self, self._selected_plan_config, ctx) + self._is_built = True + self._compiled_plans[self._plan_index].execute(self, uid_to_data, ctx) + return + + # cuDNN execution path (plan id < PYTHON_ENGINE_ID_BASE). Variant-pack + # keys are IR uids — identical to the C++ uids by construction (the IR + # owns the uid namespace and lowering pushes every uid explicitly). + from .datatypes import _is_torch_tensor + + def _ptr(d): + if type(d) is int: + return d + if _is_torch_tensor(d) or hasattr(d, "data_ptr"): + return d.data_ptr() + import cudnn + + return cudnn._pybind_module._get_data_ptr(d) # dlpack fallback + + var_pack = {uid: _ptr(d) for uid, d in uid_to_data.items()} + ws_ptr = _ptr(workspace) if workspace is not None else 0 + self._lowered_graph._execute(var_pack, ws_ptr, handle, override_uids, override_shapes, override_strides) + + def __getattr__(self, name: str): + # Plan-configuration and query methods (deselect_engines, + # get_engine_and_knobs_at_index, key, populate_cuda_graph, ...) operate + # on the lowered C++ graph — delegate to it. Only reached when normal + # attribute lookup fails, i.e. for names this class doesn't define. + lowered = self.__dict__.get("_lowered_graph") + if lowered is not None and not name.startswith("_") and hasattr(lowered, name): + return getattr(lowered, name) + raise AttributeError( + f"{type(self).__name__!r} object has no attribute {name!r}" + + ("" if lowered is not None else " (graph not lowered yet — call build_operation_graph() first)") + ) + + def __repr__(self) -> str: + if self._lowered_graph is not None: + return repr(self._lowered_graph) # classic JSON dump + import json + + return json.dumps(self.inspect(), default=str, indent=2) + + @property + def engine(self) -> Optional["BaseEngine"]: + """The python engine for the selected plan, or None for the backend path. + Populated after create_execution_plans().""" + return self.selected_engine + + def serialize(self): + """Serialize the graph (classic passthrough). + + Returns the C++ binding's serialized form unchanged; C++ + ``deserialize`` accepts exactly this form back. + """ + if self._lowered_graph is None: + # Serialization is the cuDNN graph format by definition — lower on + # demand (independent of which plan is selected for execution). + self.validate() + if self._lowered_graph is None: # python engines registered + self._lowered_graph = self._lower_to_cpp() + self._lowered_graph.validate() + self._verify_uid_ownership() + return self._lowered_graph.serialize() + + def deserialize(self, *args, **kwargs) -> None: + """Deserialize a graph (classic passthrough: (data) or (handle, data, + enforce_precompiled=...)). Replaces this graph's lowered C++ graph.""" + if self._lowered_graph is None: + import cudnn + + if self._nodes: # deserializing into a built-up graph: lower it + self.validate() + self._lowered_graph = self._lower_to_cpp() + else: # fresh container (classic usage): empty C++ graph + self._lowered_graph = cudnn._pybind_module.backend_graph() + self._lowered_graph.deserialize(*args, **kwargs) + self._is_built = True + + @classmethod + def from_serialized(cls, data, handle: Optional[int] = None, **kwargs) -> "pygraph": + """Create a pygraph from serialized data. + + This is a convenience method that creates a minimal graph and deserializes into it. + + Args: + data: Serialized graph data (from serialize()). + handle: Optional cuDNN handle for AoT compilation. + **kwargs: Additional arguments passed to the constructor. + + Returns: + pygraph: Deserialized graph ready for execution. + """ + import cudnn + + # Create a new graph with a fresh C++ graph + graph = cls(**kwargs) + graph._lowered_graph = cudnn._pybind_module.backend_graph( + io_data_type=graph._context.io_data_type, + intermediate_data_type=graph._context.intermediate_data_type, + compute_data_type=graph._context.compute_data_type, + ) + + if handle is not None: + graph._lowered_graph.deserialize(handle, data) + else: + graph._lowered_graph.deserialize(data) + graph._is_built = True + return graph + + def _lower_to_cpp(self) -> Any: + """Lower Python graph to C++ (the internal ``_pybind_module.backend_graph``).""" + import cudnn + from .datatypes import _library_type # torch dtype -> cudnn enum (classic parity) + + self._freeze() # the lowered graph mirrors the IR from here on + + # The C++ graph rejects None (wants the enum). io_data_type may be unset + # (block-scale tensors carry their own dtypes), but intermediate/compute + # default to FLOAT — matching cudnn.graph() — so cuDNN can infer virtual + # (intermediate) tensor dtypes during build. + pg_kwargs = dict(self._cpp_graph_kwargs) + if self._context.io_data_type is not None: + pg_kwargs["io_data_type"] = _library_type(self._context.io_data_type) + pg_kwargs["intermediate_data_type"] = _library_type(self._context.intermediate_data_type or cudnn.data_type.FLOAT) + pg_kwargs["compute_data_type"] = _library_type(self._context.compute_data_type or cudnn.data_type.FLOAT) + if self._handle is not None: + pg_kwargs["handle"] = self._handle + graph = cudnn._pybind_module.backend_graph(**pg_kwargs) + + tensor_map: Dict[int, Any] = {} + + def lower_tensor(t: Tensor) -> Any: + if t.uid in tensor_map: + return tensor_map[t.uid] + if t.pass_by_value is not None and t.scalar_type is not None: + cpp = graph.tensor_scalar(t.pass_by_value, t.scalar_type) + cpp.set_uid(t.uid) + tensor_map[t.uid] = cpp + return cpp + mk_kwargs = dict( + dim=t.dim, + stride=t.stride, + is_virtual=t.is_virtual, + is_pass_by_value=t.is_pass_by_value, + name=t.name, + # Always propagate the IR uid so execute()'s variant pack (keyed + # by IR uid) matches; otherwise cuDNN assigns its own and the + # buffers never bind. IR uids are unique and positive. + uid=t.uid, + ) + if t.data_type is not None: # else NOT_SET → cuDNN infers from the + mk_kwargs["data_type"] = _library_type(t.data_type) # graph intermediate default + if t.reordering_type is not None: # e.g. F8_128x4 for block-scale SFs + mk_kwargs["reordering_type"] = t.reordering_type + if t.ragged_offset is not None: + mk_kwargs["ragged_offset"] = lower_tensor(t.ragged_offset) + if t.ragged_offset_multiplier not in (None, 1): # non-default only + mk_kwargs["ragged_offset_multiplier"] = t.ragged_offset_multiplier + cpp = graph._make_tensor(**mk_kwargs) + tensor_map[t.uid] = cpp + return cpp + + def push_output_attrs(out_t: Tensor, cpp_t: Any) -> None: + # Every attribute a user may set on an OP OUTPUT via the classic + # setter chain must be pushed here (inputs get theirs through + # _make_tensor kwargs in lower_tensor). Missing one corrupts + # silently: no multiplier -> wrong GPU addresses (cudaErrorMisalignedAddress); + # no reordering -> the backend rejects or misreads the layout. + # dim/stride: USER-assigned only — inferred IR strides are + # provisional row-major; the backend keeps its classic per-op + # layout inference (channels-last conv etc.) when the user did + # not pin one. + # the label too: classic renames act on the SAME object the cpp + # graph holds, so the lowered graph carries the user's name + if out_t.name: + cpp_t.set_name(out_t.name) + if out_t.dim_assigned and out_t.dim: + cpp_t.set_dim(out_t.dim) + if out_t.stride_assigned and out_t.stride: + cpp_t.set_stride(out_t.stride) + if out_t.ragged_offset is not None: + cpp_t.set_ragged_offset(lower_tensor(out_t.ragged_offset)) + if out_t.ragged_offset_multiplier not in (None, 1): + cpp_t.set_ragged_offset_multiplier(out_t.ragged_offset_multiplier) + if out_t.reordering_type is not None: + cpp_t.set_reordering_type(out_t.reordering_type) + if not out_t.is_virtual: + cpp_t.set_output(True) + if out_t.data_type: + cpp_t.set_data_type(out_t.data_type) + + for node in self._nodes: + for t in node.inputs.values(): + if t: + lower_tensor(t) + + if node.node_type == NodeType.MATMUL: + mm_kw = dict(A=tensor_map[node.inputs["A"].uid], B=tensor_map[node.inputs["B"].uid], padding=node.params.get("padding", 0.0), name=node.name) + if node.compute_data_type is not None: + mm_kw["compute_data_type"] = _library_type(node.compute_data_type) + cpp_out = graph.matmul(**mm_kw) + elif node.node_type == NodeType.POINTWISE: + # params["mode"] IS the C++ pygraph method name — direct + # dispatch; scalar attributes (clips/negative_slope/...) are + # forwarded as keywords, tensors positionally (they lead every + # pointwise signature). + inputs = [tensor_map[t.uid] for t in node.inputs.values()] + extra = {k: node.params[k] for k in self._POINTWISE_EXTRA_PARAMS if k in node.params} + if node.compute_data_type is not None: + extra["compute_data_type"] = _library_type(node.compute_data_type) + cpp_out = getattr(graph, node.params["mode"])(*inputs, name=node.name, **extra) + elif node.node_type in _CAPTURED_BY_TYPE: + # Captured op (sdpa family): rebuild the original kwargs — + # tensor ports (port == C++ kwarg) map through tensor_map, + # scalar params forward verbatim, dropout reassembles from its + # flattened elements — and call the C++ method once. + method, spec = _CAPTURED_BY_TYPE[node.node_type] + kw = {"name": node.name} + if node.compute_data_type is not None: + kw["compute_data_type"] = _library_type(node.compute_data_type) + for pk, pv in node.params.items(): + if pk.startswith("_") or pk.startswith("dropout_"): + continue + # user callbacks (score_mod, ...) get a shimmed graph so + # closures over IR tensors keep working (see _CallbackGraphShim) + kw[pk] = _wrap_callback(pv, lower_tensor) if callable(pv) else pv + for port, t in node.inputs.items(): + if not port.startswith("dropout_"): + kw[port] = tensor_map[t.uid] + for port in spec.get("out_kwargs", ()): + if port in node.outputs: # classic passes these descriptors as args + kw[port] = lower_tensor(node.outputs[port]) + n_drop = node.params.get("_dropout_n") + if n_drop: + kw["dropout"] = tuple( + tensor_map[node.inputs[f"dropout_{i}"].uid] if f"dropout_{i}" in node.inputs else node.params[f"dropout_{i}"] for i in range(n_drop) + ) + result = getattr(graph, method)(**kw) + cpp_outs = list(result) if isinstance(result, (list, tuple)) else [result] + for oport, cpp_t in zip(spec["outputs"], cpp_outs): + out_t = node.outputs.get(oport) + if out_t is None or cpp_t is None: + continue + tensor_map[out_t.uid] = cpp_t + # sdpa-family output layout is user-chosen: the C++ node + # REQUIRES O's dim/stride before validate (BSHD vs BHSD) — + # push whatever the IR carries (inferred or user-set). + if out_t.dim: + cpp_t.set_dim(out_t.dim) + if out_t.stride: + cpp_t.set_stride(out_t.stride) + push_output_attrs(out_t, cpp_t) + continue + elif node.node_type in _STRUCTURED_BY_TYPE: + # Generic structured op (norms / reduction / block-scale / MoE / + # conv / structural): input ports are named after the C++ + # kwargs, so lowering is kwargs assembly + one call + zipping + # the returned tuple with the declared output ports. + method, spec = _STRUCTURED_BY_TYPE[node.node_type] + kw = {"name": node.name} + if not spec.get("no_cdt") and node.compute_data_type is not None: + kw["compute_data_type"] = _library_type(node.compute_data_type) + list_ports = spec.get("list_inputs", ()) + for port, t in node.inputs.items(): + if any(port.startswith(f"{lp}_") for lp in list_ports): + continue # collected below + kw[port] = tensor_map[t.uid] + for lp in list_ports: + n = node.params.get(f"_n_{lp}", 0) + if n: + kw[lp] = [tensor_map[node.inputs[f"{lp}_{i}"].uid] for i in range(n)] + for ak in spec.get("attrs", ()): + if ak in node.params: + kw[ak] = node.params[ak] + result = getattr(graph, method)(**kw) + cpp_outs = list(result) if isinstance(result, (list, tuple)) else [result] + push_dims = spec.get("push_output_dims", False) + for oport, cpp_t in zip(spec["outputs"], cpp_outs): + out_t = node.outputs.get(oport) + if out_t is None or cpp_t is None: + continue + tensor_map[out_t.uid] = cpp_t + if push_dims and out_t.dim: # ops whose output dims cuDNN can't infer + cpp_t.set_dim(out_t.dim) + # stride only when USER-assigned: pushing the IR's + # provisional row-major stride into an (e.g.) NHWC + # graph makes the backend reject the fusion (classic + # infers the stride when the user sets only dims) + if out_t.stride_assigned and out_t.stride: + cpp_t.set_stride(out_t.stride) + push_output_attrs(out_t, cpp_t) + continue + else: + continue + + # Map output + for out_t in node.outputs.values(): + tensor_map[out_t.uid] = cpp_out + push_output_attrs(out_t, cpp_out) + + # ---- uid ownership ------------------------------------------------- + # The Python IR owns the whole uid namespace: every IR tensor gets a uid + # eagerly at creation (_alloc_uid, or user-specified via tensor(uid=)), + # and lowering pushes ALL of them explicitly to C++ — inputs via + # _make_tensor(uid=), op-created outputs/virtuals via set_uid here. The + # C++ FE's build-time auto-assignment therefore NEVER triggers for + # graphs built through the Python pygraph (its enumeration order is not + # deterministic for multi-output ops, so relying on it mis-binds + # buffers). Mixed construction — adding ops directly to the lowered C++ + # graph — is unsupported: a graph is either pure-Python or pure-C++. + for ir_uid, cpp_t in tensor_map.items(): + cpp_t.set_uid(ir_uid) + + self._cpp_tensors = tensor_map + return graph + + +def _install_pointwise_builders() -> None: + """Generate the uniform pointwise builders from _POINTWISE_TENSOR_ARGS. + + Each builder accepts its tensors positionally OR by the classic pybind + keyword names (e.g. ``g.bias(input=x, bias=b)``, ``g.max(input0=a, + input1=b)``), matching the C++ pygraph API surface exactly. Ops with extra + scalar attributes (relu / leaky_relu / swish / gen_index + backwards) have + explicit builders on the class instead. + """ + + def make(op: str, argnames: tuple): + def builder(self, *args, name: str = "", compute_data_type: Any = None, **kwargs): + tensors = list(args) + for an in argnames[len(args) :]: + if an not in kwargs: + raise TypeError(f"{op}() missing tensor argument {an!r}") + tensors.append(kwargs.pop(an)) + if len(tensors) != len(argnames) or kwargs: + bad = kwargs or f"{len(tensors)} tensors" + raise TypeError(f"{op}() expects tensor arguments {argnames}; got unexpected {bad}") + return self._pointwise(op, tensors, self._get_name(op, name), compute_data_type) + + builder.__name__ = op + builder.__qualname__ = f"pygraph.{op}" + builder.__doc__ = f"Element-wise {op}({', '.join(argnames)})." + return builder + + for op, argnames in pygraph._POINTWISE_TENSOR_ARGS.items(): + if not hasattr(pygraph, op): # explicit builders (relu, ...) win + setattr(pygraph, op, make(op, argnames)) + + +_install_pointwise_builders() + + +# --------------------------------------------------------------------------- +# Structured ops, declaratively: norms, reduction, block-scale, MoE, conv, and +# the structural ops — everything except matmul (positional ergonomics) and +# sdpa (conditional kwarg assembly), which stay explicit. +# +# One table entry per op: +# node_type NodeType member (engines match on this) +# inputs ordered tensor ports == the C++ pybind kwarg names +# list_inputs ports taking a LIST of tensors (indexed ports + count) +# attrs scalar/enum/list params stored in node.params verbatim +# and forwarded as keywords at lowering +# outputs output ports, in C++ return order +# infer per-output IR-side shape inference (introspection; cuDNN +# re-infers at build) — best-effort, None on failure +# push_output_dims True for ops whose output dims cuDNN cannot infer +# (dgrad/wgrad/reduction/reshape/...): IR dims are pushed +# no_cdt True for bindings without a compute_data_type kwarg +# +# Builders are generated: tensors positionally or by port name, attrs by +# keyword, plus a reserved ``out_dims`` kwarg (dims list for a single output, +# or {port: dims} for several) for the ambiguous-shape ops. +# --------------------------------------------------------------------------- + + +def _like(port): # output dims mirror an input port + return lambda node: (node.inputs[port].dim if port in node.inputs else None) + + +def _stats_like(port, keep_axes): # input-port dims with all but keep_axes reduced to 1 + def infer(node): + d = node.inputs[port].dim if port in node.inputs else None + return [x if i in keep_axes else 1 for i, x in enumerate(d)] if d else None + + return infer + + +def _conv_fprop_dims(node): + x, w = node.inputs["image"].dim, node.inputs["weight"].dim + sp = len(x) - 2 + sym = node.params.get("padding") + pre = node.params.get("pre_padding") or sym or [0] * sp + post = node.params.get("post_padding") or sym or [0] * sp + stride = node.params.get("stride") or [1] * sp + dil = node.params.get("dilation") or [1] * sp + out = [x[0], w[0]] + for i in range(sp): + eff = (w[i + 2] - 1) * dil[i] + 1 + out.append((x[i + 2] + pre[i] + post[i] - eff) // stride[i] + 1) + return out + + +def _conv_dgrad_dims(node): + dy, w = node.inputs["loss"].dim, node.inputs["filter"].dim + sp = len(dy) - 2 + sym = node.params.get("padding") + pre = node.params.get("pre_padding") or sym or [0] * sp + post = node.params.get("post_padding") or sym or [0] * sp + stride = node.params.get("stride") or [1] * sp + dil = node.params.get("dilation") or [1] * sp + # Reverse of fprop — ambiguous for strided conv; out_dims/set_dim overrides. + out = [dy[0], w[1]] + for i in range(sp): + eff = (w[i + 2] - 1) * dil[i] + 1 + out.append((dy[i + 2] - 1) * stride[i] + eff - pre[i] - post[i]) + return out + + +def _slice_dims(node): # output extent of each python slice over the input dims + d = node.inputs["input"].dim + sls = node.params.get("slices") + if not d or not sls: + return None + return [len(range(*sl.indices(int(n)))) for sl, n in zip(sls, d)] + + +def _moe_bwd_dweight_dims(node): + do, tok, fto = (node.inputs[p].dim for p in ("doutput", "token", "first_token_offset")) + return [fto[0], tok[-1], do[-1]] # [E, H, N] + + +def _block_quant_scale_dims(node): + d = list(node.inputs["input"].dim) + bs = node.params.get("block_size") + axis = node.params.get("axis") + axis = len(d) - 1 if axis in (None, -1) else axis + d[axis] = (d[axis] + bs - 1) // bs + return d + + +_NORM_FWD_INFER = {"Y": _like("input"), "mean": _stats_like("input", (0,)), "inv_var": _stats_like("input", (0,))} +_NORM_BWD_INFER = {"DX": _like("input"), "DScale": _like("scale"), "DBias": _like("scale")} + + +def _training_phase(node): # norm stats exist only in TRAINING forward phase + phase = node.params.get("norm_forward_phase") + return getattr(phase, "name", str(phase)).upper() != "INFERENCE" + + +_NORM_FWD_MAYBE = {"mean": _training_phase, "inv_var": _training_phase} + +_STRUCTURED_OPS = { + # ---- norms -------------------------------------------------------------- + "rmsnorm": dict( + node_type=NodeType.RMSNORM, + inputs=("input", "scale", "bias", "epsilon"), + attrs=("norm_forward_phase",), + outputs=("Y", "inv_var"), + maybe={"inv_var": _training_phase}, + infer={"Y": _like("input"), "inv_var": _stats_like("input", (0,))}, + ), + "rmsnorm_backward": dict( + node_type=NodeType.RMSNORM_BWD, + inputs=("grad", "input", "scale", "inv_variance"), + attrs=("has_dbias",), + outputs=("DX", "DScale", "DBias"), + # classic rmsnorm_backward names its outputs ::Dscale/::Dbias (mixed + # case), unlike the other norm backwards (::DSCALE/::DBIAS) + out_suffix={"DScale": "Dscale", "DBias": "Dbias"}, + maybe={"DBias": lambda n: n.params.get("has_dbias", True) is not False}, + infer=_NORM_BWD_INFER, + ), + "layernorm": dict( + node_type=NodeType.LAYERNORM, + inputs=("input", "scale", "bias", "epsilon"), + attrs=("norm_forward_phase",), + outputs=("Y", "mean", "inv_var"), + maybe=_NORM_FWD_MAYBE, + infer=_NORM_FWD_INFER, + ), + "layernorm_backward": dict( + node_type=NodeType.LAYERNORM_BWD, + inputs=("grad", "input", "scale", "mean", "inv_variance"), + outputs=("DX", "DScale", "DBias"), + infer=_NORM_BWD_INFER, + ), + "adalayernorm": dict( + node_type=NodeType.ADALAYERNORM, + inputs=("input", "scale", "bias", "epsilon"), + attrs=("norm_forward_phase",), + outputs=("Y", "mean", "inv_var"), + maybe=_NORM_FWD_MAYBE, + infer=_NORM_FWD_INFER, + ), + "adalayernorm_backward": dict( + node_type=NodeType.ADALAYERNORM_BWD, + inputs=("grad", "input", "scale", "mean", "inv_variance"), + outputs=("DX", "DScale", "DBias"), + infer=_NORM_BWD_INFER, + ), + "instancenorm": dict( + node_type=NodeType.INSTANCENORM, + inputs=("input", "scale", "bias", "epsilon"), + attrs=("norm_forward_phase",), + outputs=("Y", "mean", "inv_var"), + maybe=_NORM_FWD_MAYBE, + infer={"Y": _like("input"), "mean": _stats_like("input", (0, 1)), "inv_var": _stats_like("input", (0, 1))}, + ), + "instancenorm_backward": dict( + node_type=NodeType.INSTANCENORM_BWD, + inputs=("grad", "input", "scale", "mean", "inv_variance"), + outputs=("DX", "DScale", "DBias"), + infer=_NORM_BWD_INFER, + ), + "batchnorm": dict( + node_type=NodeType.BATCHNORM, + inputs=("input", "scale", "bias", "in_running_mean", "in_running_var", "epsilon", "momentum"), + list_inputs=("peer_stats",), + outputs=("Y", "mean", "inv_var", "next_running_mean", "next_running_var"), + maybe={ + "next_running_mean": lambda n: "in_running_mean" in n.inputs, + "next_running_var": lambda n: "in_running_var" in n.inputs, + }, + infer={ + "Y": _like("input"), + "mean": _stats_like("input", (1,)), + "inv_var": _stats_like("input", (1,)), + "next_running_mean": _stats_like("input", (1,)), + "next_running_var": _stats_like("input", (1,)), + }, + ), + "batchnorm_inference": dict( + node_type=NodeType.BATCHNORM_INFERENCE, + inputs=("input", "mean", "inv_variance", "scale", "bias"), + outputs=("Y",), + infer={"Y": _like("input")}, + ), + "batchnorm_backward": dict( + node_type=NodeType.BATCHNORM_BWD, + inputs=("grad", "input", "scale", "mean", "inv_variance"), + list_inputs=("peer_stats",), + outputs=("DX", "DScale", "DBias"), + infer=_NORM_BWD_INFER, + ), + "genstats": dict( + node_type=NodeType.GENSTATS, + inputs=("input",), + outputs=("SUM", "SQ_SUM"), + infer={"SUM": _stats_like("input", (1,)), "SQ_SUM": _stats_like("input", (1,))}, + ), + # ---- reduction / block-scale / MoE -------------------------------------- + "reduction": dict( + node_type=NodeType.REDUCTION, + inputs=("input", "group_offset"), + attrs=("mode",), + outputs=("OUT_0",), + push_output_dims=True, # cuDNN needs the reduced output dims explicitly + ), + "block_scale_dequantize": dict( + node_type=NodeType.BLOCK_SCALE_DEQUANTIZE, + inputs=("input", "descale"), + attrs=("block_size", "is_negative_scale"), + outputs=("OUT_0",), + infer={"OUT_0": _like("input")}, + ), + "block_scale_quantize": dict( + node_type=NodeType.BLOCK_SCALE_QUANTIZE, + inputs=("input",), + attrs=("block_size", "axis", "transpose"), + outputs=("Y", "scale"), + infer={"Y": _like("input"), "scale": _block_quant_scale_dims}, + ), + "moe_grouped_matmul": dict( + node_type=NodeType.MOE_GROUPED_MATMUL, + inputs=("token", "weight", "first_token_offset", "token_index", "token_ks"), + attrs=("mode", "top_k"), + outputs=("OUT_0",), + infer={"OUT_0": lambda n: [1, n.inputs["token"].dim[-2], n.inputs["weight"].dim[-1]]}, + ), + "moe_grouped_matmul_bwd": dict( + node_type=NodeType.MOE_GROUPED_MATMUL_BWD, + inputs=("doutput", "token", "first_token_offset"), + outputs=("dweight",), + infer={"dweight": _moe_bwd_dweight_dims}, + push_output_dims=True, + ), + # ---- convolution --------------------------------------------------------- + "conv_fprop": dict( + node_type=NodeType.CONV_FPROP, + inputs=("image", "weight"), + attrs=("padding", "pre_padding", "post_padding", "stride", "dilation", "convolution_mode"), + outputs=("Y",), + infer={"Y": _conv_fprop_dims}, + ), + "conv_dgrad": dict( + node_type=NodeType.CONV_DGRAD, + inputs=("loss", "filter"), + attrs=("padding", "pre_padding", "post_padding", "stride", "dilation", "convolution_mode"), + outputs=("DX",), + infer={"DX": _conv_dgrad_dims}, + push_output_dims=True, # dgrad output dims are ambiguous for strided conv + ), + "conv_wgrad": dict( + node_type=NodeType.CONV_WGRAD, + inputs=("image", "loss"), + attrs=("padding", "pre_padding", "post_padding", "stride", "dilation", "convolution_mode"), + outputs=("DW",), + push_output_dims=True, # wgrad output (filter) dims are not inferable + ), + # ---- structural ----------------------------------------------------------- + "reshape": dict( + node_type=NodeType.RESHAPE, + inputs=("input",), + attrs=("reshape_mode",), + outputs=("OUT_0",), + push_output_dims=True, # target shape comes from out_dims / set_dim + no_cdt=True, + ), + "slice": dict( + node_type=NodeType.SLICE, + inputs=("input",), + attrs=("slices",), + outputs=("OUT_0",), + infer={"OUT_0": _slice_dims}, + dtype_like={"OUT_0": "input"}, # classic: slice output dtype == input's + ), + "transpose": dict( + node_type=NodeType.TRANSPOSE, + inputs=("input",), + attrs=("permutation",), + outputs=("OUT_0",), + infer={"OUT_0": lambda n: ([n.inputs["input"].dim[i] for i in n.params["permutation"]] if n.inputs["input"].dim else None)}, + ), + "concatenate": dict( + node_type=NodeType.CONCATENATE, + inputs=(), + list_inputs=("inputs",), + attrs=("axis", "in_place_index"), + outputs=("OUT_0",), + no_cdt=True, + ), + "rope": dict( + node_type=NodeType.ROPE, + inputs=("input", "freqs"), + attrs=("output_scale", "rope_dim"), + outputs=("OUT_0",), + infer={"OUT_0": _like("input")}, + ), + "rope_backward": dict( + node_type=NodeType.ROPE_BWD, + inputs=("dY", "freqs"), + attrs=("output_scale", "rope_dim"), + outputs=("OUT_0",), + infer={"OUT_0": _like("dY")}, + ), +} + +# node_type -> (method name, spec), for the generic lowering branch +_STRUCTURED_BY_TYPE = {spec["node_type"]: (op, spec) for op, spec in _STRUCTURED_OPS.items()} + + +def _install_structured_builders() -> None: + """Generate builders for _STRUCTURED_OPS. + + Call style: tensors positionally (in declared port order) or by port name; + attrs by keyword; ``out_dims`` sets output dims explicitly (a dims list for + single-output ops, or {port: dims}) for shapes cuDNN cannot infer.""" + + def make(op: str, spec: dict): + input_ports = spec["inputs"] + list_ports = spec.get("list_inputs", ()) + attr_kws = spec.get("attrs", ()) + infer = spec.get("infer", {}) + maybe = spec.get("maybe", {}) + + def builder(self, *args, name: str = "", compute_data_type: Any = None, out_dims: Any = None, **kwargs): + name_ = self._get_name(op, name) + node = Node(name_, spec["node_type"], compute_data_type or self._context.compute_data_type) + # classic positional order: tensor ports first, then attrs + n_p = len(input_ports) + if len(args) > n_p + len(attr_kws): + raise TypeError(f"{op}() takes at most {n_p + len(attr_kws)} positional arguments ({input_ports} + {attr_kws})") + for ak, v in zip(attr_kws, args[n_p:]): + if ak in kwargs: + raise TypeError(f"{op}() got multiple values for {ak!r}") + kwargs[ak] = v + args = args[:n_p] + for port, v in zip(input_ports, args): + node.inputs[port] = self._ensure_tensor(v, name=f"{name_}::{port}") + for port in input_ports[len(args) :]: + v = kwargs.pop(port, None) + if v is not None: + node.inputs[port] = self._ensure_tensor(v, name=f"{name_}::{port}") + for lp in list_ports: + vs = kwargs.pop(lp, None) or [] + for i, v in enumerate(vs): + node.inputs[f"{lp}_{i}"] = self._ensure_tensor(v, name=f"{name_}::{lp}_{i}") + if vs: + node.params[f"_n_{lp}"] = len(vs) + for ak in attr_kws: + v = kwargs.pop(ak, None) + if v is not None: + node.params[ak] = v + if kwargs: + raise TypeError(f"{op}() got unexpected arguments {sorted(kwargs)}; tensor ports are {input_ports}, attrs are {attr_kws}") + if out_dims is not None and not isinstance(out_dims, dict): + out_dims = {spec["outputs"][0]: out_dims} + dtype_like = spec.get("dtype_like", {}) + outs = [] + for oport in spec["outputs"]: + cond = maybe.get(oport) + if cond is not None and not cond(node): + outs.append(None) # classic returns None for absent outputs + continue + o = self._make_output(f"{name_}::{spec.get('out_suffix', {}).get(oport) or self._CLASSIC_OUT_SUFFIX.get(oport, oport)}") + src = dtype_like.get(oport) + if src and src in node.inputs: + o.data_type = node.inputs[src].data_type + d = (out_dims or {}).get(oport) + if d is None: + try: # best-effort IR-side inference; C++ validates at build + d = infer.get(oport, lambda n: None)(node) + except Exception: # noqa: BLE001 + d = None + if d: + o.dim = list(d) + o.stride = _row_major_stride(o.dim) + node.outputs[oport] = o + self._register_tensor(o) + outs.append(o) + self._nodes.append(node) + return outs[0] if len(outs) == 1 else list(outs) # classic multi-output ops return a LIST + + builder.__name__ = op + builder.__qualname__ = f"pygraph.{op}" + builder.__doc__ = f"{op}({', '.join(input_ports)}) -> ({', '.join(spec['outputs'])})." + return builder + + for op, spec in _STRUCTURED_OPS.items(): + setattr(pygraph, op, make(op, spec)) + + +_install_structured_builders() + + +# --------------------------------------------------------------------------- +# Captured ops (the sdpa family): the kwarg surface is huge (~130 args across +# the six variants, including tensor-or-float args, dropout tuples, and +# score_mod callbacks), so builders capture ALL kwargs generically instead of +# hand-mirroring each one — tensors become named ports (port == C++ kwarg), +# everything else goes to params verbatim, dropout tuples are flattened per +# element. Lowering rebuilds the kwargs and makes one C++ call. The node stays +# first-class: engines read node.inputs["q"] / node.params["use_causal_mask"]. +# --------------------------------------------------------------------------- + + +class _CallbackGraphShim: + """Wraps the C++ graph handed to user callbacks (score_mod & co.) during + lowering. Classic code passes the SAME object at build and callback time, so + closures over user-created tensors just work; post-flip the user's closures + capture IR Tensors while the callback receives the C++ graph. The shim + translates any IR Tensor argument to its lowered C++ tensor at the call + site, so existing callback code runs unchanged.""" + + def __init__(self, target, lower_tensor): + self._target = target + self._lower = lower_tensor + + def _xlate(self, v): + if isinstance(v, Tensor): + # closure-captured helper tensors may not feed any node: lower on demand + return self._lower(v) + if isinstance(v, (list, tuple)): + return type(v)(self._xlate(x) for x in v) + return v + + def __getattr__(self, name): + attr = getattr(self._target, name) + if not callable(attr): + return attr + + def call(*args, **kwargs): + return attr(*[self._xlate(a) for a in args], **{k: self._xlate(v) for k, v in kwargs.items()}) + + return call + + +def _wrap_callback(fn, lower_tensor): + """Wrap a user callback param (e.g. score_mod) so the C++ graph it receives + is shimmed (see _CallbackGraphShim) and stray IR-Tensor args translate.""" + import functools + + @functools.wraps(fn) + def wrapped(*args, **kwargs): + import cudnn + + cpp_graph_t = cudnn._pybind_module.backend_graph + + def conv(v): + if isinstance(v, cpp_graph_t): + return _CallbackGraphShim(v, lower_tensor) + if isinstance(v, Tensor): + return lower_tensor(v) + return v + + return fn(*[conv(a) for a in args], **{k: conv(v) for k, v in kwargs.items()}) + + return wrapped + + +def _stats_expected(params): + if params.get("generate_stats") is not None: + return bool(params["generate_stats"]) + return not params.get("is_inference", True) + + +def _sdpa_o_dims(node): # O: q dims with v's head dim + q, v = node.inputs["q"].dim, node.inputs["v"].dim + return list(q[:-1]) + [v[-1]] + + +def _sdpa_stats_dims(node): # Stats: q dims with last dim 1 + return list(node.inputs["q"].dim[:-1]) + [1] + + +_AMAX = lambda node: [1, 1, 1, 1] # noqa: E731 — fp8 amax side outputs + +_CAPTURED_OPS = { + "sdpa": dict( + node_type=NodeType.SDPA, + pos=("q", "k", "v"), + outputs=("O", "Stats"), + # kwargs whose tensors are semantically OUTPUTS of the node (the classic + # API passes their descriptors as arguments): recorded in node.outputs so + # engines see correct producer/consumer direction. + out_kwargs=("rng_dump", "score_max", "score_sum_exp"), + maybe={"Stats": _stats_expected}, + infer={"O": _sdpa_o_dims, "Stats": _sdpa_stats_dims}, + ), + "sdpa_backward": dict( + node_type=NodeType.SDPA_BWD, + pos=("q", "k", "v", "o", "dO", "stats"), + outputs=("dQ", "dK", "dV"), + out_kwargs=("dBias", "dSink_token", "rng_dump"), + infer={"dQ": _like("q"), "dK": _like("k"), "dV": _like("v")}, + ), + "sdpa_fp8": dict( + node_type=NodeType.SDPA_FP8, + pos=("q", "k", "v", "descale_q", "descale_k", "descale_v", "descale_s", "scale_s", "scale_o"), + outputs=("O", "Stats", "Amax_S", "Amax_O"), + out_kwargs=("rng_dump", "score_max", "score_sum_exp"), + maybe={"Stats": _stats_expected}, + infer={"O": _sdpa_o_dims, "Stats": _sdpa_stats_dims, "Amax_S": _AMAX, "Amax_O": _AMAX}, + ), + "sdpa_fp8_backward": dict( + node_type=NodeType.SDPA_FP8_BWD, + pos=( + "q", + "k", + "v", + "o", + "dO", + "stats", + "descale_q", + "descale_k", + "descale_v", + "descale_o", + "descale_dO", + "descale_s", + "descale_dP", + "scale_s", + "scale_dQ", + "scale_dK", + "scale_dV", + "scale_dP", + ), + outputs=("dQ", "dK", "dV", "amax_dQ", "amax_dK", "amax_dV", "amax_dP"), + out_kwargs=("dSink_token",), + infer={"dQ": _like("q"), "dK": _like("k"), "dV": _like("v"), "amax_dQ": _AMAX, "amax_dK": _AMAX, "amax_dV": _AMAX, "amax_dP": _AMAX}, + ), + # mxfp8 variants (schemas match the bindings exactly; output dims via + # out_dims / set_dim where cuDNN needs them) + "sdpa_mxfp8": dict( + node_type=NodeType.SDPA_MXFP8, + pos=("q", "k", "v", "descale_q", "descale_k", "descale_v"), + outputs=("O", "Stats", "Amax_O"), + maybe={"Stats": _stats_expected}, + ), + "sdpa_mxfp8_backward": dict( + node_type=NodeType.SDPA_MXFP8_BWD, + pos=( + "q", + "q_T", + "k", + "k_T", + "v", + "o_f16", + "dO_f16", + "dO", + "dO_T", + "stats", + "descale_q", + "descale_q_T", + "descale_k", + "descale_k_T", + "descale_v", + "descale_dO", + "descale_dO_T", + ), + outputs=("dQ", "dK", "dV", "amax_dQ", "amax_dK", "amax_dV"), + out_kwargs=("dSink_token",), + ), +} + +_CAPTURED_BY_TYPE = {spec["node_type"]: (op, spec) for op, spec in _CAPTURED_OPS.items()} + + +def _install_captured_builders() -> None: + """Generate the sdpa-family builders (generic kwarg capture).""" + + def _tensorish(v): + return isinstance(v, Tensor) or hasattr(v, "__dlpack__") + + def make(op: str, spec: dict): + pos = spec.get("pos", ()) + infer = spec.get("infer", {}) + maybe = spec.get("maybe", {}) + + def builder(self, *args, name: str = "", compute_data_type: Any = None, out_dims: Any = None, **kwargs): + name_ = self._get_name(op, name) + node = Node(name_, spec["node_type"], compute_data_type or self._context.compute_data_type) + if len(args) > len(pos): + raise TypeError(f"{op}() takes at most {len(pos)} positional arguments {pos}") + for k, v in zip(pos, args): + if k in kwargs: + raise TypeError(f"{op}() got multiple values for {k!r}") + kwargs[k] = v + drop = kwargs.pop("dropout", None) + out_kwargs = spec.get("out_kwargs", ()) + for k, v in kwargs.items(): + if v is None: + continue + if _tensorish(v): + if k in out_kwargs: # semantically an OUTPUT of this node + node.outputs[k] = self._ensure_tensor(v, name=f"{name_}::{k}") + else: + node.inputs[k] = self._ensure_tensor(v, name=f"{name_}::{k}") + else: # scalar / enum / callback — forwarded verbatim at lowering + node.params[k] = v + if drop is not None: + node.params["_dropout_n"] = len(drop) + for i, e in enumerate(drop): + if _tensorish(e): + node.inputs[f"dropout_{i}"] = self._ensure_tensor(e, name=f"{name_}::dropout_{i}") + else: + node.params[f"dropout_{i}"] = e + if out_dims is not None and not isinstance(out_dims, dict): + out_dims = {spec["outputs"][0]: out_dims} + rets = [] + for oport in spec["outputs"]: + cond = maybe.get(oport) + if cond is not None and not cond(node.params): + rets.append(None) # e.g. Stats in inference mode (classic returns None) + continue + o = self._make_output(f"{name_}::{spec.get('out_suffix', {}).get(oport) or self._CLASSIC_OUT_SUFFIX.get(oport, oport)}") + d = (out_dims or {}).get(oport) + if d is None: + try: + d = infer.get(oport, lambda n: None)(node) + except Exception: # noqa: BLE001 + d = None + if d: + o.dim = list(d) + o.stride = _row_major_stride(o.dim) + node.outputs[oport] = o + self._register_tensor(o) + rets.append(o) + self._nodes.append(node) + return list(rets) # always full arity; classic returns a LIST + + builder.__name__ = op + builder.__qualname__ = f"pygraph.{op}" + builder.__doc__ = f"{op}(...) -> {spec['outputs']} (generic kwarg capture; see _CAPTURED_OPS)." + return builder + + for op, spec in _CAPTURED_OPS.items(): + setattr(pygraph, op, make(op, spec)) + + +_install_captured_builders() diff --git a/python/cudnn/engines/__init__.py b/python/cudnn/engines/__init__.py new file mode 100644 index 000000000..253bae4b2 --- /dev/null +++ b/python/cudnn/engines/__init__.py @@ -0,0 +1,30 @@ +"""Execution backends for pygraph. + +Pluggable execution backends in one flat engine-id space with the cuDNN backend. +The Router builds a ranked plan list at ``create_execution_plans()`` time; graph +construction stays backend-agnostic. + +Backends: +- ReferenceMatmulEngine: pure-PyTorch correctness oracle (CPU/GPU, no JIT deps) + +Real DSL engines (cuTile / CuTe-DSL GEMM fusion) plug in as separate PRs — an +engine is one file implementing BaseEngine; nothing here changes. +""" + +from .base import BaseEngine, CompiledPlan, ExecutionContext, PlanConfig +from .engine_ids import PYTHON_ENGINE_ID_BASE, BACKEND_HEURISTIC_ENGINE_ID, is_python_engine +from .router import Router, default_router +from .reference_matmul_engine import ReferenceMatmulEngine + +__all__ = [ + "BaseEngine", + "Router", + "PlanConfig", + "CompiledPlan", + "ExecutionContext", + "default_router", + "ReferenceMatmulEngine", + "PYTHON_ENGINE_ID_BASE", + "BACKEND_HEURISTIC_ENGINE_ID", + "is_python_engine", +] diff --git a/python/cudnn/engines/base.py b/python/cudnn/engines/base.py new file mode 100644 index 000000000..bed12beff --- /dev/null +++ b/python/cudnn/engines/base.py @@ -0,0 +1,167 @@ +"""Backend (engine) contract for the Python graph: plan -> compile -> execute. + +A backend is one of the interchangeable implementations the Router dispatches +to (Python DSLs, a naive reference, the cuDNN Graph backend, ...). The +lifecycle mirrors a real JIT/DSL engine: + + 1. ``propose_plans(graph)`` -> candidate ``PlanConfig`` entries (one per + configuration the engine wants ranked; decline the whole graph by raising + ``NotImplementedError`` / ``cudnn.cudnnGraphNotSupportedError``). + 2. ``build_plan(graph, plan)`` -> a ``CompiledPlan`` — the expensive JIT step, + run ONCE per (graph, selected plan) at ``graph.build_plans()`` time; the + compiled artifact lives on the graph, so one engine instance is safely + reusable across graphs. + 3. ``CompiledPlan.execute(graph, tensor_data, ctx)`` — hot path. The + ``ExecutionContext`` carries the caller's handle / stream / workspace / + dynamic-shape overrides explicitly; engines must not hard-code a stream or + silently allocate hidden workspace. + +Simple eager engines only implement ``execute()`` — the default ``build_plan`` +wraps it in a trivial ``CompiledPlan``. + +Example: + class MyEngine(BaseEngine): + name = "my_engine" + engine_id = PYTHON_ENGINE_ID_BASE + 7 # stable id it owns + + def check_support(self, graph): + for node in graph.nodes: + if node.node_type != NodeType.MATMUL: + raise NotImplementedError(...) + + def execute(self, graph, tensor_data, ctx): + ... # write results into caller-provided output buffers +""" + +from abc import ABC +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Dict, List + +from .engine_ids import PYTHON_ENGINE_ID_BASE # noqa: F401 — re-exported for engine authors + +if TYPE_CHECKING: + from ..pygraph import pygraph + + +@dataclass(frozen=True) +class PlanConfig: + """One candidate execution plan: an engine id + its knobs. + + ``engine_id`` lives in the shared flat id space (``engine_ids``); knobs are + engine-specific tuning (cuDNN knob dict, or a python engine's config). The + plan's source is derived from the id via ``is_python_engine`` — no separate + field, so cuDNN and python plans are interchangeable in the ranked list. + One engine may propose several plans differing only in knobs. + """ + + engine_id: int + knobs: Any = None + + +@dataclass(frozen=True) +class ExecutionContext: + """Runtime context passed to a compiled plan at execute time. + + Everything an engine may need is explicit here — no engine should reach + into private graph state, hard-code a stream, or allocate hidden workspace. + ``stream`` is resolved from the handle when available (classic + ``cudnn.set_stream(handle, ...)`` semantics). + """ + + handle: Any = None + stream: Any = None + workspace: Any = None + override_uids: Any = None + override_shapes: Any = None + override_strides: Any = None + + +class CompiledPlan: + """A compiled (graph, plan) artifact. Subclass for real JIT engines.""" + + def get_workspace_size(self) -> int: + """Workspace bytes this plan needs at execute time (default 0).""" + return 0 + + def execute(self, graph: "pygraph", tensor_data: Dict[int, Any], ctx: ExecutionContext) -> None: + raise NotImplementedError + + +class _EagerPlan(CompiledPlan): + """Default CompiledPlan for simple eager engines (delegates to engine.execute).""" + + def __init__(self, engine: "BaseEngine", plan: PlanConfig): + self.engine = engine + self.plan = plan + + def get_workspace_size(self) -> int: + return self.engine.get_workspace_size() + + def execute(self, graph, tensor_data, ctx: ExecutionContext) -> None: + self.engine.execute(graph, tensor_data, ctx) + + +class BaseEngine(ABC): + """Abstract base class for python graph execution backends. + + Attributes: + name: Human-readable identifier. + engine_id: Stable id in the shared flat engine-id space, in the reserved + Python region (>= PYTHON_ENGINE_ID_BASE). Subclasses MUST declare it; + the base default (None) is rejected at register_backend(). + default_knobs: Optional default tuning knobs for this engine's plan. + """ + + name: str = "base" + # Subclasses MUST declare a stable id in the reserved python region; the + # base intentionally has none so a forgotten override fails at registration + # instead of silently colliding with another engine. + engine_id: Any = None + default_knobs: Any = None + + def __init__(self): + pass + + def check_support(self, graph: "pygraph") -> None: + """Raise to decline ``graph``. + + Decline ONLY via ``NotImplementedError`` or + ``cudnn.cudnnGraphNotSupportedError`` (the classic unsupported-graph + signal); any other exception is treated as an engine bug and propagates. + Default: accept everything (subclasses should narrow this). + """ + _ = graph + + def propose_plans(self, graph: "pygraph") -> List[PlanConfig]: + """Candidate plans for ``graph``, in this engine's preference order. + + Default: one plan with ``default_knobs`` when ``check_support`` accepts. + Engines with several viable configurations override this to expose them + to ranking/autotune (each entry's knobs reach ``build_plan`` verbatim). + """ + self.check_support(graph) + return [PlanConfig(self.engine_id, self.default_knobs)] + + def build_plan(self, graph: "pygraph", plan: PlanConfig, ctx: "ExecutionContext" = None) -> CompiledPlan: + """Compile ``graph`` for ``plan`` (the expensive step; run once per + graph/plan at build_plans() time). ``ctx`` carries the build context — + handle and stream when available — so device-specific AoT compilers get + their inputs explicitly instead of reading private graph state. Default + wraps eager ``execute()``.""" + return _EagerPlan(self, plan) + + def get_workspace_size(self) -> int: + """Workspace bytes for eager engines (compiled plans report their own).""" + return 0 + + def execute(self, graph: "pygraph", tensor_data: Dict[int, Any], ctx: ExecutionContext) -> None: + """Eager execution hook (used by the default ``build_plan``). + + ``tensor_data`` maps tensor UIDs (inputs + outputs) to their device + data. Write results directly into the caller-provided output buffers, + on ``ctx.stream`` when set. + """ + raise NotImplementedError(f"Engine '{self.name}' must implement execute() or build_plan()") + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(name={self.name!r}, engine_id={self.engine_id})" diff --git a/python/cudnn/engines/engine_ids.py b/python/cudnn/engines/engine_ids.py new file mode 100644 index 000000000..2920fc1be --- /dev/null +++ b/python/cudnn/engines/engine_ids.py @@ -0,0 +1,32 @@ +"""Engine-id namespace shared by cuDNN and Python backends. + +Execution engines live in one flat integer id space, exactly like cuDNN's own +backend engines (which have small ids 0..N, each with knobs). Python engines +occupy a reserved high region so the two never collide and a heuristics query +can return a single ranked list mixing both, e.g.: + + [(engine_id=1048576, knobs), (engine_id=1, knobs), (engine_id=5, knobs), ...] + +Dispatch is a single predicate on the id: ``is_python_engine(id)`` -> run via the +Python engine registry; otherwise lower to the cuDNN C++ backend. + +Each Python engine declares a *stable* ``engine_id`` in this range (it owns its +id, the way a cuDNN engine does), so ids don't shift with registration order — +autotune results and pinned plans stay reproducible across runs. +""" + +# Start of the reserved Python-engine id region. 1<<20 (~1.05M) is far above any +# plausible cuDNN engine count, so the two id spaces can never collide without +# having to know cuDNN's actual maximum. +PYTHON_ENGINE_ID_BASE = 1 << 20 + +# The backend side of the plan list: "delegate to the loaded backend's own +# heuristics". Deliberately ONE entry — the backend's engine set varies by +# backend version and is only discoverable per graph at plan time, never +# statically enumerable by the frontend. +BACKEND_HEURISTIC_ENGINE_ID = -1 + + +def is_python_engine(engine_id: int) -> bool: + """True iff ``engine_id`` names a Python engine (vs a cuDNN backend engine).""" + return engine_id >= PYTHON_ENGINE_ID_BASE diff --git a/python/cudnn/engines/reference_matmul_engine.py b/python/cudnn/engines/reference_matmul_engine.py new file mode 100644 index 000000000..daff71ca9 --- /dev/null +++ b/python/cudnn/engines/reference_matmul_engine.py @@ -0,0 +1,93 @@ +"""Pure-PyTorch reference backend — a correctness baseline with no GPU/JIT deps. + +This backend exists so the pygraph + BaseEngine + Router contract can be +exercised in CI on CPU, and so every future DSL backend has a numerical oracle +to diff against. It supports MATMUL plus a small set of POINTWISE ops; anything +else is declined (the Router then tries another backend or falls back to cuDNN). + +It runs wherever the input tensors live (CPU or CUDA) via ``torch.matmul`` / +elementwise ops, and writes results into the caller-provided output buffers. +""" + +from typing import TYPE_CHECKING, Any, Dict + +try: + import torch +except ImportError: + torch = None + +from .base import BaseEngine +from .engine_ids import PYTHON_ENGINE_ID_BASE +from ..graph_types import NodeType + +if TYPE_CHECKING: + from ..pygraph import pygraph + +# POINTWISE ops this reference understands, keyed by the op kind +# (params["mode"] == the pygraph method name). +_UNARY = { + "relu": lambda x: x.clamp_min(0), + "gelu": lambda x: torch.nn.functional.gelu(x), + "sigmoid": lambda x: torch.sigmoid(x), + "tanh": lambda x: torch.tanh(x), + "exp": lambda x: torch.exp(x), + "identity": lambda x: x, +} +_BINARY = { + "add": lambda a, b: a + b, + "mul": lambda a, b: a * b, + "sub": lambda a, b: a - b, + "div": lambda a, b: a / b, + "bias": lambda a, b: a + b, + "scale": lambda a, b: a * b, +} + + +def _mode_name(mode: Any) -> str: + return getattr(mode, "name", str(mode)).lower() + + +class ReferenceMatmulEngine(BaseEngine): + """CPU/GPU PyTorch reference for MATMUL + basic POINTWISE fusions.""" + + name = "reference_matmul" + engine_id = PYTHON_ENGINE_ID_BASE + 0 # stable id (a correctness oracle) + + def check_support(self, graph: "pygraph") -> None: + if torch is None: + raise NotImplementedError("ReferenceMatmulEngine requires PyTorch") + for node in graph.nodes: + if node.node_type == NodeType.MATMUL: + continue + if node.node_type == NodeType.POINTWISE: + mode = _mode_name(node.params.get("mode")) + if mode not in _UNARY and mode not in _BINARY: + raise NotImplementedError(f"ReferenceMatmulEngine: unsupported pointwise mode {mode!r}") + if any(k != "mode" for k in node.params): + # scalar attributes (clips / negative_slope / ...) not implemented + raise NotImplementedError("ReferenceMatmulEngine: pointwise scalar attributes not supported") + continue + raise NotImplementedError(f"ReferenceMatmulEngine only supports MATMUL / basic POINTWISE, got {node.node_type.name}") + + def execute(self, graph, tensor_data: Dict[int, Any], ctx=None) -> None: + # Nodes are already in build (topological) order. Compute each node into + # a scratch map, then copy declared outputs into the caller's buffers. + values: Dict[int, Any] = dict(tensor_data) + + for node in graph.nodes: + if node.node_type == NodeType.MATMUL: + a = values[node.inputs["A"].uid] + b = values[node.inputs["B"].uid] + out = torch.matmul(a, b) + elif node.node_type == NodeType.POINTWISE: + mode = _mode_name(node.params.get("mode")) + ins = [values[t.uid] for t in node.inputs.values()] + out = _UNARY[mode](ins[0]) if mode in _UNARY else _BINARY[mode](ins[0], ins[1]) + else: # pragma: no cover — guarded by check_support + raise NotImplementedError(node.node_type.name) + + out_t = next(iter(node.outputs.values())) + dst = values.get(out_t.uid) + if dst is not None and hasattr(dst, "copy_"): + dst.copy_(out) # caller-provided output buffer + values[out_t.uid] = out diff --git a/python/cudnn/engines/router.py b/python/cudnn/engines/router.py new file mode 100644 index 000000000..f3b20970e --- /dev/null +++ b/python/cudnn/engines/router.py @@ -0,0 +1,93 @@ +"""Router: builds the ranked execution-plan list at plan-creation time. + +Implements the dispatch stage of the Python API unification proposal: + + Python Graph API -> create_execution_plans() -> Router -> ranked plan list + (one flat (engine_id, + knobs) list mixing + python DSLs + backend) + +Routing happens at ``create_execution_plans()`` time, NOT at graph construction, +so graph building stays backend-agnostic (lazy lowering). The Router returns a +flat list of ``PlanConfig(engine_id, knobs)``: Python engines (ids in the +reserved high region) whose ``check_support()`` accepts the graph, plus AT MOST +ONE backend delegating entry (``BACKEND_HEURISTIC_ENGINE_ID``). Dispatch on each +plan's id (``is_python_engine``) decides whether to run via the Python registry +or lower to the cuDNN C++ backend. + +WHAT THIS MR SUPPORTS (the enforced contract — ``create_execution_plans()`` +validates the final Router output, whatever the Router implementation): + +* python entries must name engines registered on the graph; +* the only legal non-python entry is ONE backend delegating sentinel — the + backend's own plans stay behind it, addressed via the classic at-index APIs + (a separate, backend-owned index space); +* an empty plan list is rejected (there is no legal empty planning state). + +Concrete backend engine configs as first-class routed entries +(``PlanConfig(cudnn_engine_id, knobs)`` interleaved with python plans) are NOT +representable in this MR: they need a typed plan representation and a build +path via cpp ``create_execution_plan(engine_id, knobs)`` — that is the +heuristics/autotune follow-up MR's job, not one extra lowering branch. What IS +already decided and stable here: routed indices never shift (the sentinel never +expands in place), and the backend engine set is discovered per graph at plan +time (get_engine_count / get_engine_and_knobs_at_index on the lowered graph) — +never statically enumerated in frontend code, because it varies by backend +version. + +Policy remains pluggable at three levels: subclass ``Router`` and override +``plan()``; pass per-graph via ``pygraph(router=...)`` / ``set_router()`` +(before planning); or swap the process-wide ``default_router``. ``plan()`` may +return any ordering/mix of the representable entries — python-first, +backend-first, interleaved, conditional on the graph. The current default is a +placeholder concat. +""" + +from typing import TYPE_CHECKING, List + +from .base import BaseEngine, PlanConfig +from .engine_ids import BACKEND_HEURISTIC_ENGINE_ID + +if TYPE_CHECKING: + from .._pygraph import pygraph + + +class Router: + """Default policy: python engines that support the graph, then the backend.""" + + def plan(self, graph: "pygraph", backends: List[BaseEngine]) -> List[PlanConfig]: + """Return the ranked candidate plan list for ``graph``. + + Python engines are included (by ascending ``engine_id``, a stable order) + when their ``check_support(graph)`` does not raise. A backend DECLINES + only via ``NotImplementedError`` or ``cudnn.cudnnGraphNotSupportedError`` + (the classic unsupported-graph signal); any other exception is a bug in + the engine and propagates to the caller instead of silently falling back. + """ + import cudnn + + decline = (NotImplementedError, cudnn.cudnnGraphNotSupportedError) + plans: List[PlanConfig] = [] + for engine in sorted(backends, key=lambda e: e.engine_id): + try: + proposals = engine.propose_plans(graph) + except decline: + continue + for pc in proposals: + if pc.engine_id != engine.engine_id: # no identity injection + raise ValueError(f"engine {engine.name!r} proposed a plan with foreign engine_id {pc.engine_id}") + plans.extend(proposals) + + # The backend side is ONE delegating entry by design: the frontend owns + # only its python-engine id segment and must work against any (incl. + # future) backend version, so the backend's engine set can never be + # statically enumerated here — it is discovered per graph at plan time + # via the backend's own heuristics/query API (get_engine_and_knobs_at_ + # index on the lowered graph) when a caller wants to expand or autotune. + plans.append(PlanConfig(BACKEND_HEURISTIC_ENGINE_ID)) + return plans + + +# Process-wide default. Assign a Router subclass to change global policy, or pass +# one to pygraph(router=...) / graph.set_router(...) per graph. +default_router = Router() diff --git a/python/cudnn/graph_types.py b/python/cudnn/graph_types.py new file mode 100644 index 000000000..f08806791 --- /dev/null +++ b/python/cudnn/graph_types.py @@ -0,0 +1,234 @@ +"""Pure Python data types for cuDNN Frontend graph representation. + +This module provides Python dataclasses for tensor attributes and node types, +enabling native Python access to graph structure. +""" + +from dataclasses import dataclass, field +from enum import Enum, auto +from typing import Any, Dict, List, Optional, Union + + +class NodeType(Enum): + """Operation node types. Maps to INode::Type in node_interface.h.""" + + # Only the op types exercised by this version are listed. Add more as + # needed, following the block-scale / MoE / reduction examples (enum entry + # here + a builder in graph_native + inference in nodes + lowering). + COMPOSITE = auto() + CONV_FPROP = auto() + CONV_DGRAD = auto() + CONV_WGRAD = auto() + GENSTATS = auto() + RESHAPE = auto() + SLICE = auto() + CONCATENATE = auto() + TRANSPOSE = auto() + ROPE = auto() + ROPE_BWD = auto() + MOE_GROUPED_MATMUL_BWD = auto() + MATMUL = auto() + MATMUL_FP8 = auto() + POINTWISE = auto() + REDUCTION = auto() + RMSNORM = auto() + RMSNORM_BWD = auto() + LAYERNORM = auto() + LAYERNORM_BWD = auto() + ADALAYERNORM = auto() + ADALAYERNORM_BWD = auto() + INSTANCENORM = auto() + INSTANCENORM_BWD = auto() + BATCHNORM = auto() + BATCHNORM_INFERENCE = auto() + BATCHNORM_BWD = auto() + SDPA = auto() + SDPA_BWD = auto() + SDPA_FP8 = auto() + SDPA_FP8_BWD = auto() + SDPA_MXFP8 = auto() + SDPA_MXFP8_BWD = auto() + MOE_GROUPED_MATMUL = auto() + BLOCK_SCALE_QUANTIZE = auto() + BLOCK_SCALE_DEQUANTIZE = auto() + + +@dataclass(eq=False) # identity-based hash/eq: uid/name are mutable +class Tensor: + """Pure Python representation of tensor attributes. + + Mirrors cudnn_frontend::graph::Tensor_attributes from graph_properties.h. + + Attributes: + name: Tensor identifier + data_type: Data type (uses cudnn.data_type values) + dim: Dimensions of the tensor + stride: Memory strides + is_virtual: True if tensor is an intermediate (not I/O) + is_pass_by_value: True if tensor is a scalar passed at execution + pass_by_value: Embedded constant value (for fused scalars) + uid: Unique identifier for backend mapping + uid_assigned: True if UID was explicitly assigned + reordering_type: Memory layout transformation type + ragged_offset: Tensor for variable-length tensor offsets + """ + + name: str = "" + data_type: Any = None + dim: List[int] = field(default_factory=list) + stride: List[int] = field(default_factory=list) + is_virtual: bool = False + is_pass_by_value: bool = False + pass_by_value: Optional[Union[int, float]] = None + uid: int = 0 + uid_assigned: bool = False + # user-assigned vs IR-inferred layout: only USER-assigned dim/stride are + # pushed to the lowered C++ output tensors — inferred values are + # provisional (row-major) and the backend applies its own classic + # per-op layout inference (e.g. channels-last conv). Internal inference + # writes the attributes directly and leaves these False. + dim_assigned: bool = False + stride_assigned: bool = False + reordering_type: Any = None + ragged_offset: Optional["Tensor"] = None + ragged_offset_multiplier: int = 1 + scalar_type: Any = None # cudnn.scalar_type for tensor_scalar-created scalars + # weakref to the owning graph (set at registration): identity mutations + # (set_name / set_uid) delegate to the graph so its indexes stay coherent. + owner: Any = field(default=None, repr=False) + + def __setattr__(self, name, value): + # direct attribute writes freeze with the owning graph (the fluent + # setters are guarded separately and give a richer error) + if getattr(self, "_frozen", False) and name != "_frozen": + raise RuntimeError(f"cannot set Tensor.{name}: the owning graph is frozen after lowering/planning") + object.__setattr__(self, name, value) + + def _guard(self, what: str = "mutate a tensor attribute") -> None: + g = self.owner() if self.owner is not None else None + if g is not None: + g._check_mutable(what) + + def set_output(self, value: bool) -> "Tensor": + """Mark this tensor as an output (non-virtual) or intermediate (virtual).""" + self._guard() + self.is_virtual = not value + return self + + def set_data_type(self, dtype: Any) -> "Tensor": + """Set the data type.""" + self._guard() + self.data_type = dtype + return self + + def set_name(self, name: str) -> "Tensor": + """Set the tensor name (graph-owned tensors re-index atomically).""" + g = self.owner() if self.owner is not None else None + if g is not None: + g._rename_tensor(self, name) + else: + self.name = name + return self + + def set_dim(self, dim: List[int]) -> "Tensor": + """Set the tensor dimensions (user-assigned: pushed at lowering).""" + self._guard() + self.dim = list(dim) + self.dim_assigned = True + return self + + def set_stride(self, stride: List[int]) -> "Tensor": + """Set the tensor strides (user-assigned: pushed at lowering).""" + self._guard() + self.stride = list(stride) + self.stride_assigned = True + return self + + def set_uid(self, uid: int) -> "Tensor": + """Set the tensor UID (graph-owned tensors re-index atomically; a + colliding auto-assigned uid is renumbered, user-user conflicts raise).""" + g = self.owner() if self.owner is not None else None + if g is not None: + g._reuid_tensor(self, uid) + else: + self.uid = uid + self.uid_assigned = True + return self + + def set_ragged_offset(self, ragged_offset: "Tensor") -> "Tensor": + """Set the ragged-offset tensor (variable-length layouts).""" + self._guard() + self.ragged_offset = ragged_offset + return self + + def set_ragged_offset_multiplier(self, multiplier: int) -> "Tensor": + """Set the ragged-offset unit size in tensor elements.""" + self._guard() + self.ragged_offset_multiplier = multiplier + return self + + def set_reordering_type(self, reordering_type: Any) -> "Tensor": + """Set the memory reordering layout (e.g. F8_128x4).""" + self._guard() + self.reordering_type = reordering_type + return self + + def set_is_pass_by_value(self, value: bool) -> "Tensor": + """Mark the tensor as a host pass-by-value scalar.""" + self._guard() + self.is_pass_by_value = value + return self + + def set_is_virtual(self, value: bool) -> "Tensor": + """Set virtualness directly (classic parity; inverse of set_output).""" + self._guard() + self.is_virtual = value + return self + + def get_uid(self) -> int: + return self.uid + + def get_name(self) -> str: + return self.name + + def get_dim(self) -> List[int]: + return list(self.dim) # a copy, like the classic pybind getter + + def get_stride(self) -> List[int]: + return list(self.stride) # a copy, like the classic pybind getter + + def get_data_type(self) -> Any: + # classic parity: the pybind getter returns the cudnn enum even when + # the user set a torch dtype (classic converts at set time) + dt = self.data_type + if dt is not None and type(dt).__module__ == "torch": + from .datatypes import _torch_to_cudnn_data_type + + return _torch_to_cudnn_data_type(dt) + return self.data_type + + def get_is_virtual(self) -> bool: + return self.is_virtual + + def get_is_pass_by_value(self) -> bool: + return self.is_pass_by_value + + def get_reordering_type(self) -> Any: + return self.reordering_type + + def get_ragged_offset_multiplier(self) -> int: + return self.ragged_offset_multiplier + + def validate(self) -> None: + """Validate tensor configuration.""" + if not self.dim: + raise ValueError(f"Tensor '{self.name}' dims not set.") + if not self.stride: + raise ValueError(f"Tensor '{self.name}' strides not set.") + if len(self.dim) != len(self.stride): + raise ValueError(f"Tensor '{self.name}' dim/stride length mismatch: " f"{len(self.dim)} vs {len(self.stride)}") + if self.is_virtual and self.is_pass_by_value: + raise ValueError(f"Tensor '{self.name}' can't be both virtual and pass_by_value.") + + # NOTE: hash/eq are object identity (dataclass eq=False). uid and name are + # mutable, so value-based hashing would violate the dict-key invariant. diff --git a/python/cudnn/nodes.py b/python/cudnn/nodes.py new file mode 100644 index 000000000..8a3537142 --- /dev/null +++ b/python/cudnn/nodes.py @@ -0,0 +1,163 @@ +"""Pure Python node class for cuDNN Frontend graph representation. + +Simple, Pythonic design - everything stored directly on the Node. +""" + +from typing import TYPE_CHECKING, Dict, List, Optional, Any + +from .graph_types import NodeType, Tensor + +if TYPE_CHECKING: + from ._pygraph import GraphContext + + +class Node: + """A single operation node in the computation graph. + + All operation-specific parameters are stored in the `params` dict, + keeping the design simple and flexible. + + Attributes: + name: Operation name + node_type: Type of operation (MATMUL, POINTWISE, SDPA, etc.) + inputs: Dict mapping port names to input Tensor + outputs: Dict mapping port names to output Tensor + params: Dict of operation-specific parameters (padding, stride, mode, etc.) + compute_data_type: Data type for computation + """ + + def __init__( + self, + name: str, + node_type: NodeType, + compute_data_type: Any = None, + ): + self.name = name + self.node_type = node_type + self.compute_data_type = compute_data_type + self.inputs: Dict[str, Tensor] = {} + self.outputs: Dict[str, Tensor] = {} + self.params: Dict[str, Any] = {} + + def __setattr__(self, name, value): + # attribute writes freeze with the owning graph; the port/param dicts + # themselves become MappingProxy views at freeze time + if getattr(self, "_frozen", False) and name != "_frozen": + raise RuntimeError(f"cannot set Node.{name}: the owning graph is frozen after lowering/planning") + object.__setattr__(self, name, value) + + def validate(self) -> None: + """Validate node configuration.""" + for port_name, tensor in self.inputs.items(): + if tensor is None: + raise ValueError(f"Node '{self.name}': Input '{port_name}' is None") + for port_name, tensor in self.outputs.items(): + if tensor is None: + raise ValueError(f"Node '{self.name}': Output '{port_name}' is None") + + if self.node_type == NodeType.MATMUL: + self._validate_matmul() + + def infer_properties(self, context: "GraphContext") -> None: + """Infer unset tensor properties from context and inputs.""" + # Fill data types from context + for tensor in self.inputs.values(): + if tensor and tensor.data_type is None: + tensor.data_type = context.intermediate_data_type if tensor.is_virtual else context.io_data_type + + for tensor in self.outputs.values(): + if tensor and tensor.data_type is None: + tensor.data_type = context.intermediate_data_type if tensor.is_virtual else context.io_data_type + + # Operation-specific inference + if self.node_type == NodeType.MATMUL: + self._infer_matmul() + elif self.node_type == NodeType.POINTWISE: + self._infer_pointwise() + # structured/captured ops (norms/conv/moe/sdpa/...) infer at build time + # via their tables' per-output lambdas + + def _validate_matmul(self) -> None: + """Validate matmul dimensions: C = A @ B.""" + a = self.inputs.get("A") + b = self.inputs.get("B") + if not (a and b and a.dim and b.dim): + return + if a.dim[-1] != b.dim[-2]: + raise ValueError(f"Node '{self.name}': Inner dimensions must match for matmul: " f"A{a.dim} @ B{b.dim}") + + def _infer_matmul(self) -> None: + """Infer output dims for matmul: C = A @ B.""" + a = self.inputs.get("A") + b = self.inputs.get("B") + c = self.outputs.get("C") + + if not (a and b and c): + return + + if not c.dim and a.dim and b.dim: + # Output shape: [..., M, N] where M=A[-2], N=B[-1] + ndim = max(len(a.dim), len(b.dim)) + c_dim = [1] * ndim + + # Last two dims: M from A, N from B + if len(a.dim) >= 2: + c_dim[-2] = a.dim[-2] + if len(b.dim) >= 2: + c_dim[-1] = b.dim[-1] + + # Broadcast batch dims (incompatible extents raise, matching numpy + # rules: equal, or one side is 1) + for i in range(ndim - 2): + a_idx = i - (ndim - len(a.dim)) + b_idx = i - (ndim - len(b.dim)) + a_val = a.dim[a_idx] if 0 <= a_idx < len(a.dim) - 2 else 1 + b_val = b.dim[b_idx] if 0 <= b_idx < len(b.dim) - 2 else 1 + if a_val != b_val and 1 not in (a_val, b_val): + raise ValueError(f"Node '{self.name}': batch dims not broadcastable: A{a.dim} vs B{b.dim}") + c_dim[i] = max(a_val, b_val) + + c.dim = c_dim + + if not c.stride and c.dim: + c.stride = _row_major_stride(c.dim) + + def _infer_pointwise(self) -> None: + """Infer output dims for pointwise (broadcast inputs).""" + out = self.outputs.get("OUT_0") + if not out: + return + + if not out.dim: + # Right-aligned elementwise broadcast across all inputs (numpy + # rules); lower-rank operands contribute to the trailing dims. + max_dim: list = [] + for tensor in self.inputs.values(): + if not (tensor and tensor.dim): + continue + d = list(tensor.dim) + if len(d) > len(max_dim): + d, max_dim = max_dim, d # keep max_dim the longer one + for i in range(1, len(d) + 1): # merge right-aligned + a, b = max_dim[-i], d[-i] + if a != b and 1 not in (a, b): + raise ValueError(f"Node '{self.name}': pointwise inputs not broadcastable") + max_dim[-i] = max(a, b) + if max_dim: + out.dim = max_dim + + if not out.stride and out.dim: + out.stride = _row_major_stride(out.dim) + + def __repr__(self) -> str: + return f"Node({self.name!r}, {self.node_type.name})" + + +def _row_major_stride(dim: List[int]) -> List[int]: + """Compute row-major (C-contiguous) strides.""" + if not dim: + return [] + stride = [1] * len(dim) + for i in range(len(dim) - 2, -1, -1): + stride[i] = stride[i + 1] * dim[i + 1] + return stride diff --git a/python/cudnn/wrapper.py b/python/cudnn/wrapper.py index 950a07fb4..56d696c55 100644 --- a/python/cudnn/wrapper.py +++ b/python/cudnn/wrapper.py @@ -37,6 +37,12 @@ from typing import Any, Dict, List, Optional, Tuple, Union import cudnn + +# Graph tensors come in two forms post-unification: the classic pybind +# ``cudnn.tensor`` (deserialized/legacy graphs) and the Python-IR +# ``cudnn.Tensor`` returned by ``cudnn.pygraph`` ops. Both duck-type the same +# getters (get_name/get_uid/get_dim/...). +_GRAPH_TENSOR_TYPES = (cudnn.tensor, cudnn.Tensor) import cudnn.datatypes from cudnn import data_type, heur_mode @@ -112,7 +118,7 @@ def _find_tensor( for tensor_name, tensor_value in tensor_map.items(): if tensor_value.get_uid() == tensor: return tensor_name - elif isinstance(tensor, cudnn.tensor): + elif isinstance(tensor, _GRAPH_TENSOR_TYPES): for tensor_name, tensor_value in tensor_map.items(): if tensor is tensor_value: return tensor_name @@ -398,7 +404,7 @@ def wrapper(*args, **kwargs): if obj_id not in self.__tensor_map: self.__tensor_map[obj_id] = _graph_tensor(self.__graph, elem) obj[j] = self.__tensor_map[obj_id] - if isinstance(obj[j], cudnn.tensor): + if isinstance(obj[j], _GRAPH_TENSOR_TYPES): self.__tensor_in[f"{node_name}::{i}::{j}"] = obj[j] obj = args[i] = tuple(obj) # convert back to tuple if hasattr(obj, "__dlpack__"): @@ -406,7 +412,7 @@ def wrapper(*args, **kwargs): if obj_id not in self.__tensor_map: self.__tensor_map[obj_id] = _graph_tensor(self.__graph, obj) obj = args[i] = self.__tensor_map[obj_id] - if isinstance(obj, cudnn.tensor): + if isinstance(obj, _GRAPH_TENSOR_TYPES): self.__tensor_in[f"{node_name}::{i}"] = obj # process keyword arguments for dlpack tensors for key, obj in kwargs.items(): @@ -419,7 +425,7 @@ def wrapper(*args, **kwargs): if obj_id not in self.__tensor_map: self.__tensor_map[obj_id] = _graph_tensor(self.__graph, elem) obj[j] = self.__tensor_map[obj_id] - if isinstance(obj[j], cudnn.tensor): + if isinstance(obj[j], _GRAPH_TENSOR_TYPES): self.__tensor_in[f"{node_name}::{key}::{j}"] = obj[j] obj = kwargs[key] = tuple(obj) # convert back to tuple if hasattr(obj, "__dlpack__"): @@ -427,16 +433,16 @@ def wrapper(*args, **kwargs): if obj_id not in self.__tensor_map: self.__tensor_map[obj_id] = _graph_tensor(self.__graph, obj) obj = kwargs[key] = self.__tensor_map[obj_id] - if isinstance(obj, cudnn.tensor): + if isinstance(obj, _GRAPH_TENSOR_TYPES): self.__tensor_in[f"{node_name}::{key}"] = obj # capturing node output output = attr(*args, **kwargs) - if isinstance(output, cudnn.tensor): + if isinstance(output, _GRAPH_TENSOR_TYPES): output_list = [output] elif isinstance(output, (list, tuple)): output_list = output for i, obj in enumerate(output_list): - if isinstance(obj, cudnn.tensor): + if isinstance(obj, _GRAPH_TENSOR_TYPES): if hasattr(obj, "get_name") and obj.get_name(): tensor_name = obj.get_name() else: diff --git a/python/pygraph/pygraph.cpp b/python/pygraph/pygraph.cpp index fb6bf37aa..ed04f9b19 100644 --- a/python/pygraph/pygraph.cpp +++ b/python/pygraph/pygraph.cpp @@ -829,7 +829,7 @@ default_vector(void) { void init_pygraph_submodule(py::module_& m) { - py::class_ pygraph_(m, "pygraph"); + py::class_ pygraph_(m, "backend_graph"); pygraph_ .def(py::init 0 and not u.uid_assigned # -1 == auto + assert u.data_type is None or getattr(u.data_type, "name", "") != "NOT_SET" diff --git a/test/python/test_engine_router.py b/test/python/test_engine_router.py new file mode 100644 index 000000000..ef23a68b0 --- /dev/null +++ b/test/python/test_engine_router.py @@ -0,0 +1,431 @@ +"""CPU tests for the backend Router + BaseEngine contract. + +These run without a GPU or cuDNN: they exercise pygraph -> Router -> ranked +plan list -> engine-id dispatch using the pure-PyTorch ReferenceMatmulEngine. +This is the CI-safe proof that the unification contract works end to end. +""" + +import pytest + +torch = pytest.importorskip("torch") + +from cudnn._pygraph import pygraph +from cudnn.engines import BaseEngine, Router, ReferenceMatmulEngine, PYTHON_ENGINE_ID_BASE, is_python_engine +from cudnn.engines.engine_ids import BACKEND_HEURISTIC_ENGINE_ID + +pytestmark = pytest.mark.L0 + + +def test_router_plan_list_includes_supporting_engines_then_cudnn(): + """Plan list = supporting python engines (by id) + a trailing backend entry.""" + + class Declines(BaseEngine): + name = "declines" + engine_id = PYTHON_ENGINE_ID_BASE + 50 + + def check_support(self, graph): + raise NotImplementedError("nope") + + def execute(self, graph, tensor_data, ctx=None): + raise AssertionError("should not run") + + class Accepts(BaseEngine): + name = "accepts" + engine_id = PYTHON_ENGINE_ID_BASE + 10 + + def execute(self, graph, tensor_data, ctx=None): + pass + + g = pygraph() + a = g.tensor(dim=[4, 8], name="A") + b = g.tensor(dim=[8, 4], name="B") + g.matmul(a, b, name="mm") + g.register_backend(Accepts()).register_backend(Declines()) + + plans = Router().plan(g, g.backends) + ids = [p.engine_id for p in plans] + # Only the supporting python engine is included, then the backend entry last. + assert ids == [PYTHON_ENGINE_ID_BASE + 10, BACKEND_HEURISTIC_ENGINE_ID] + assert is_python_engine(ids[0]) and not is_python_engine(ids[-1]) + + +def test_reference_matmul_execute_cpu(): + """ReferenceMatmulEngine runs a matmul on CPU and writes the output buffer.""" + g = pygraph() + g.register_backend(ReferenceMatmulEngine()) + + a = torch.randn(2, 3, 4) + b = torch.randn(2, 4, 5) + C = g.matmul(a, b, name="mm") + c = torch.empty(2, 3, 5) + + g.execute({C: c}) + + assert g.selected_engine is not None + assert g.selected_engine.name == "reference_matmul" + torch.testing.assert_close(c, torch.matmul(a, b)) + + +def test_reference_matmul_bias_relu_fusion_cpu(): + """A small matmul + add + relu chain routes to the reference and matches.""" + g = pygraph() + g.register_backend(ReferenceMatmulEngine()) + + a = torch.randn(3, 4) + b = torch.randn(4, 5) + bias = torch.randn(3, 5) + mm = g.matmul(a, b, name="mm") + biased = g.add(mm, g._ensure_tensor(bias, name="bias"), name="bias_add") + out = g.relu(biased, name="act") + c = torch.empty(3, 5) + + g.execute({out: c}) + + ref = torch.relu(torch.matmul(a, b) + bias) + torch.testing.assert_close(c, ref) + + +def test_select_plan_survives_build_and_execute(): + """Regression (review item 2): select_plan(i) must not be reset by the + implicit build() inside execute().""" + + class EngA(BaseEngine): + name = "a" + engine_id = PYTHON_ENGINE_ID_BASE + 10 + ran = 0 + + def execute(self, graph, tensor_data, ctx=None): + type(self).ran += 1 + + class EngB(BaseEngine): + name = "b" + engine_id = PYTHON_ENGINE_ID_BASE + 11 + ran = 0 + + def execute(self, graph, tensor_data, ctx=None): + type(self).ran += 1 + + g = pygraph() + g.register_backend(EngA()).register_backend(EngB()) + a = torch.randn(2, 2) + C = g.matmul(a, torch.randn(2, 2)) + g.create_execution_plans() + assert g.selected_engine.name == "a" + g.select_plan(1) # pin engine B + assert g.selected_engine.name == "b" + g.execute({C: torch.empty(2, 2)}) # implicit build() must preserve the pin + assert EngB.ran == 1 and EngA.ran == 0 + + +def test_register_backend_validation(): + """Regression (review item 6): duplicate/invalid ids rejected at + registration; registration after planning rejected.""" + + class NoId(BaseEngine): + name = "noid" # forgets to declare engine_id (base default is None) + + def execute(self, graph, tensor_data, ctx=None): + pass + + class E1(BaseEngine): + name = "e1" + engine_id = PYTHON_ENGINE_ID_BASE + 20 + + def execute(self, graph, tensor_data, ctx=None): + pass + + g = pygraph() + with pytest.raises(ValueError, match="engine_id"): + g.register_backend(NoId()) + g.register_backend(E1()) + with pytest.raises(ValueError, match="already registered"): + g.register_backend(E1()) + a = g.tensor(dim=[2, 2], name="A") + g.matmul(a, g.tensor(dim=[2, 2], name="B")) + g.create_execution_plans() + with pytest.raises(RuntimeError, match="after create_execution_plans"): + + class E2(E1): + engine_id = PYTHON_ENGINE_ID_BASE + 21 + + g.register_backend(E2()) + + +def test_unexpected_engine_exception_propagates(): + """Regression (review item 6): only NotImplementedError / + cudnnGraphNotSupportedError decline; other exceptions are engine bugs.""" + + class Buggy(BaseEngine): + name = "buggy" + engine_id = PYTHON_ENGINE_ID_BASE + 30 + + def check_support(self, graph): + raise RuntimeError("driver exploded") + + def execute(self, graph, tensor_data, ctx=None): + pass + + g = pygraph() + a = g.tensor(dim=[2, 2], name="A") + g.matmul(a, g.tensor(dim=[2, 2], name="B")) + g.register_backend(Buggy()) + with pytest.raises(RuntimeError, match="driver exploded"): + g.create_execution_plans() + + +def test_no_backend_plan_list_is_cudnn_only(): + """With no python engine, the plan list is just the backend entry (selected=None).""" + g = pygraph() + a = g.tensor(dim=[4, 8], name="A") + b = g.tensor(dim=[8, 4], name="B") + g.matmul(a, b, name="mm") + g.validate() + + from cudnn.engines.router import default_router + + plans = default_router.plan(g, g.backends) + assert [p.engine_id for p in plans] == [BACKEND_HEURISTIC_ENGINE_ID] + g._plans = plans + assert g.selected_engine is None # backend path + + +def test_compiled_plan_lifecycle_knobs_and_reuse(): + """Review item 1 acceptance: multiple knob proposals from one engine; the + selected plan's knobs reach build_plan; compilation runs once per plan and + the artifact is reused; caller workspace + stream context reach execute.""" + from cudnn.engines import CompiledPlan, ExecutionContext, PlanConfig + + compiled_log = [] + + class TunablePlan(CompiledPlan): + def __init__(self, knobs): + self.knobs = knobs + self.executed = [] + + def get_workspace_size(self): + return 4096 + + def execute(self, graph, tensor_data, ctx): + self.executed.append((self.knobs, ctx.workspace)) + + class Tunable(BaseEngine): + name = "tunable" + engine_id = PYTHON_ENGINE_ID_BASE + 40 + + def propose_plans(self, graph): + return [PlanConfig(self.engine_id, {"tile": 128}), PlanConfig(self.engine_id, {"tile": 256})] + + def build_plan(self, graph, plan, ctx=None): + compiled_log.append(plan.knobs) + return TunablePlan(plan.knobs) + + g = pygraph() + g.register_backend(Tunable()) + C = g.matmul(torch.randn(2, 2), torch.randn(2, 2)) + g.create_execution_plans() + assert [p.knobs for p in g.plans[:2]] == [{"tile": 128}, {"tile": 256}] + + ws = torch.empty(4096, dtype=torch.uint8) + out = torch.empty(2, 2) + g.select_plan(1) # the tile=256 plan + assert len(g.plans) == 3 # two knob proposals + the backend delegating entry + g.build_plans() + assert compiled_log == [{"tile": 256}] # compiled once, correct knobs + assert g.get_workspace_size() == 4096 # plan-specific workspace + g.execute({C: out}, workspace=ws) + g.execute({C: out}, workspace=ws) + assert compiled_log == [{"tile": 256}] # reused, no recompilation + plan = g._compiled_plans[g._plan_index] + assert plan.executed[0] == ({"tile": 256}, ws) # knobs + caller workspace observed + + # same engine instance on a second graph: no state collision + g2 = pygraph() + g2.register_backend(Tunable()) + C2 = g2.matmul(torch.randn(2, 2), torch.randn(2, 2)) + g2.execute({C2: torch.empty(2, 2)}) + assert compiled_log == [{"tile": 256}, {"tile": 128}] # g2 compiled its own plan + + +def _mk_engine(id_off, knobs=None, log=None): + from cudnn.engines import CompiledPlan, PlanConfig + + class _Plan(CompiledPlan): + def __init__(self, k): + self.knobs = k + + def execute(self, graph, tensor_data, ctx): + (log if log is not None else []).append(self.knobs) + + class _E(BaseEngine): + name = f"e{id_off}" + engine_id = PYTHON_ENGINE_ID_BASE + id_off + default_knobs = knobs + + def build_plan(self, graph, plan, ctx=None): + return _Plan(plan.knobs) + + def execute(self, graph, tensor_data, ctx=None): + pass + + return _E() + + +def test_planning_is_one_shot(): + """Classic conformance: re-planning was never a supported call pattern (the + C++ graph appends plans by accident on a second call; nobody re-plans). + A second create_execution_plans() raises — a stale compiled artifact can + therefore never execute. Plan differently => build a new graph.""" + log = [] + eng = _mk_engine(60, knobs="old", log=log) + g = pygraph() + g.register_backend(eng) + C = g.matmul(torch.randn(2, 2), torch.randn(2, 2)) + g.create_execution_plans() + g.build_plans() + with pytest.raises(RuntimeError, match="one-shot"): + g.create_execution_plans() + g.execute({C: torch.empty(2, 2)}) + assert log[-1] == "old" # the planned artifact, unchanged + + +def test_mixed_router_ordering_dispatch(): + """Follow-up item 2: dispatch honors arbitrary Router ordering (backend-first, + interleaved), never a python-prefix assumption.""" + from cudnn.engines import PlanConfig, Router + from cudnn.engines.engine_ids import BACKEND_HEURISTIC_ENGINE_ID + + ran = [] + ea, eb = _mk_engine(61, "A", ran), _mk_engine(62, "B", ran) + + class Interleaved(Router): + def plan(self, graph, backends): + return [ + PlanConfig(ea.engine_id, "A"), + PlanConfig(BACKEND_HEURISTIC_ENGINE_ID), + PlanConfig(eb.engine_id, "B"), + ] + + g = pygraph(router=Interleaved()) + g.register_backend(ea).register_backend(eb) + C = g.matmul(torch.randn(2, 2), torch.randn(2, 2)) + g.create_execution_plans() + # slot 0 = python A, slot 1 = cuDNN, slot 2 = python B + assert g.selected_engine.name == "e61" + g.select_plan(1) # the backend delegating entry is selectable in place + assert g.selected_engine is None # None == the backend path + g.select_plan(2) + assert g.selected_engine.name == "e62" + g.execute({C: torch.empty(2, 2)}) + assert ran[-1] == "B" + # the middle routed entry is the backend delegating one, and routed indices + # are STABLE: python-B stays at index 2 regardless of lowering. (Real + # execution THROUGH the backend slot of a mixed router is the GPU test + # test_mixed_router_backend_slot_executes in test_native_backend_lowering.py.) + assert g.plans[1].engine_id == BACKEND_HEURISTIC_ENGINE_ID + assert g.selected_engine.name == "e62" + + +def test_empty_router_output_rejected(): + """A Router returning [] is an error — there is no legal empty planning + state (it would defeat the one-shot flag and every needs-planning check).""" + + class Empty(Router): + def plan(self, graph, backends): + return [] + + g = pygraph(router=Empty()) + g.matmul(torch.randn(2, 2), torch.randn(2, 2)) + with pytest.raises(ValueError, match="empty plan list"): + g.create_execution_plans() + # the failed call did NOT consume the one-shot: fixing the router by + # rebuilding the graph is the documented path, but the graph must not be + # left half-planned either + assert not g._planning_done + + +def test_set_router_frozen_after_planning(): + """set_router() after planning raises (it could not affect the already + planned list; accepting it silently would lie).""" + g = pygraph() + g.register_backend(_mk_engine(70)) + g.matmul(torch.randn(2, 2), torch.randn(2, 2)) + g.create_execution_plans() + with pytest.raises(RuntimeError, match="one-shot"): + g.set_router(Router()) + + +def test_backend_count_is_a_separate_space(): + """get_execution_plan_count() is the classic backend-count passthrough, + never the routed-list length; the routed list is graph.plans/select_plan. + A python-only routed graph has no backend plans and says so.""" + + class PythonOnly(Router): + def plan(self, graph, backends): + from cudnn.engines import PlanConfig + + return [PlanConfig(backends[0].engine_id)] + + g = pygraph(router=PythonOnly()) + g.register_backend(_mk_engine(71)) + C = g.matmul(torch.randn(2, 2), torch.randn(2, 2)) + g.create_execution_plans() + assert len(g.plans) == 1 + with pytest.raises(RuntimeError, match="graph.plans"): + g.get_execution_plan_count() # no backend entry -> no backend plans + g.execute({C: torch.empty(2, 2)}) # the routed python plan still runs + + +def test_constructor_backends_validated_and_proposals_checked(): + """Follow-up item 6: constructor path uses registration validation; foreign + engine ids in proposals are rejected.""" + from cudnn.engines import PlanConfig + + class NoId(BaseEngine): + def execute(self, graph, tensor_data, ctx=None): + pass + + with pytest.raises(ValueError, match="engine_id"): + pygraph(backends=[NoId()]) + + class Impostor(BaseEngine): + name = "impostor" + engine_id = PYTHON_ENGINE_ID_BASE + 63 + + def propose_plans(self, graph): + return [PlanConfig(PYTHON_ENGINE_ID_BASE + 99, None)] # foreign id + + def execute(self, graph, tensor_data, ctx=None): + pass + + g = pygraph(backends=[Impostor()]) + g.matmul(torch.randn(2, 2), torch.randn(2, 2)) + with pytest.raises(ValueError, match="foreign engine_id"): + g.create_execution_plans() + + +def test_python_plan_rejects_dynamic_workspace_overrides(): + g = pygraph() + g.register_backend(ReferenceMatmulEngine()) + C = g.matmul(torch.randn(2, 3), torch.randn(3, 2)) + g.build() + with pytest.raises(NotImplementedError, match="overrides"): + g.get_workspace_size(1234) + g.execute({C: torch.empty(2, 2)}) # normal path unaffected + + +def test_failed_stream_query_on_supplied_handle_raises(monkeypatch): + """Follow-up item 3: a supplied handle whose stream cannot be queried is a + correctness error — never a silent stream-0 fallback.""" + import cudnn as _cudnn + + g = pygraph() + g.register_backend(ReferenceMatmulEngine()) + C = g.matmul(torch.randn(2, 3), torch.randn(3, 2)) + g.build() + + def boom(handle): + raise RuntimeError("stream query failed") + + monkeypatch.setattr(_cudnn, "get_stream", boom) + with pytest.raises(RuntimeError, match="stream query failed"): + g.execute({C: torch.empty(2, 2)}, handle=42) diff --git a/test/python/test_graph_native.py b/test/python/test_graph_native.py new file mode 100644 index 000000000..603eacb05 --- /dev/null +++ b/test/python/test_graph_native.py @@ -0,0 +1,648 @@ +"""Unit tests for Python-native graph representation.""" + +import pytest + +torch = pytest.importorskip("torch") + +from cudnn.graph_types import NodeType, Tensor +from cudnn.nodes import Node, _row_major_stride +from cudnn._pygraph import pygraph, GraphContext + +pytestmark = pytest.mark.L0 + + +class TestTensor: + """Tests for Tensor.""" + + def test_create_tensor(self): + t = Tensor(name="test", dim=[8, 64, 128], stride=[8192, 128, 1]) + assert t.name == "test" + assert t.dim == [8, 64, 128] + assert not t.is_virtual + + def test_builder_pattern(self): + t = Tensor() + t.set_name("my_tensor").set_dim([4, 32]).set_stride([32, 1]).set_output(True) + assert t.name == "my_tensor" + assert not t.is_virtual + + def test_set_output(self): + t = Tensor(is_virtual=True) + t.set_output(True) + assert not t.is_virtual + t.set_output(False) + assert t.is_virtual + + def test_validation_success(self): + t = Tensor(name="valid", dim=[8, 64], stride=[64, 1]) + t.validate() + + def test_validation_no_dims(self): + t = Tensor(name="no_dims", stride=[64, 1]) + with pytest.raises(ValueError, match="dims not set"): + t.validate() + + def test_validation_dim_stride_mismatch(self): + t = Tensor(name="mismatch", dim=[8, 64, 128], stride=[64, 1]) + with pytest.raises(ValueError, match="mismatch"): + t.validate() + + def test_uid_management(self): + t = Tensor(name="test") + assert not t.uid_assigned + t.set_uid(42) + assert t.uid == 42 + assert t.uid_assigned + + +class TestNode: + """Tests for Node class.""" + + def test_create_node(self): + node = Node("mm1", NodeType.MATMUL) + assert node.name == "mm1" + assert node.node_type == NodeType.MATMUL + assert node.inputs == {} + assert node.outputs == {} + assert node.params == {} + + def test_node_with_tensors(self): + node = Node("mm1", NodeType.MATMUL) + a = Tensor(name="A", dim=[8, 64], stride=[64, 1]) + b = Tensor(name="B", dim=[64, 32], stride=[32, 1]) + c = Tensor(name="C", dim=[8, 32], stride=[32, 1]) + + node.inputs["A"] = a + node.inputs["B"] = b + node.outputs["C"] = c + + assert node.inputs["A"] is a + assert node.outputs["C"] is c + + def test_node_params(self): + node = Node("mm1", NodeType.MATMUL) + node.params["padding"] = 0.0 + node.params["alpha"] = 2.0 + assert node.params["padding"] == 0.0 + + def test_node_repr(self): + node = Node("mm1", NodeType.MATMUL) + assert repr(node) == "Node('mm1', MATMUL)" + + +class TestMatmulInference: + """Tests for matmul dimension inference.""" + + def test_infer_2d(self): + node = Node("mm", NodeType.MATMUL) + a = Tensor(name="A", dim=[64, 128], stride=[128, 1]) + b = Tensor(name="B", dim=[128, 256], stride=[256, 1]) + c = Tensor(name="C", is_virtual=True) + + node.inputs["A"] = a + node.inputs["B"] = b + node.outputs["C"] = c + + node.infer_properties(GraphContext()) + assert c.dim == [64, 256] + + def test_infer_3d_batched(self): + node = Node("mm", NodeType.MATMUL) + a = Tensor(name="A", dim=[8, 64, 128], stride=[8192, 128, 1]) + b = Tensor(name="B", dim=[8, 128, 256], stride=[32768, 256, 1]) + c = Tensor(name="C", is_virtual=True) + + node.inputs["A"] = a + node.inputs["B"] = b + node.outputs["C"] = c + + node.infer_properties(GraphContext()) + assert c.dim == [8, 64, 256] + + def test_infer_strides(self): + node = Node("mm", NodeType.MATMUL) + a = Tensor(name="A", dim=[8, 64], stride=[64, 1]) + b = Tensor(name="B", dim=[64, 32], stride=[32, 1]) + c = Tensor(name="C", is_virtual=True) + + node.inputs["A"] = a + node.inputs["B"] = b + node.outputs["C"] = c + + node.infer_properties(GraphContext()) + assert c.stride == [32, 1] + + +class TestRowMajorStride: + """Tests for stride computation.""" + + def test_1d(self): + assert _row_major_stride([10]) == [1] + + def test_2d(self): + assert _row_major_stride([8, 64]) == [64, 1] + + def test_3d(self): + assert _row_major_stride([4, 8, 16]) == [128, 16, 1] + + def test_empty(self): + assert _row_major_stride([]) == [] + + +class Testpygraph: + """Tests for pygraph.""" + + def test_creation(self): + g = pygraph() + assert len(g.nodes) == 0 + assert len(g.tensors) == 0 + + def test_with_context(self): + g = pygraph(io_data_type="HALF", compute_data_type="FLOAT") + assert g.context.io_data_type == "HALF" + assert g.context.compute_data_type == "FLOAT" + + def test_tensor_creation(self): + g = pygraph() + t = g.tensor(dim=[8, 64, 128], name="my_tensor") + assert t.name == "my_tensor" + assert t.dim == [8, 64, 128] + assert t.stride == [8192, 128, 1] + assert "my_tensor" in g.tensors + + def test_uid_ownership(self): + """The IR owns the uid namespace: user-specified uids are reserved (auto + allocation skips them). The SAME collision rule applies at creation as + at set_uid: a user uid landing on an auto-assigned one steals it (the + auto holder is renumbered — classic tensors have no uid until assigned, + so classic code cannot observe auto uids); user-user collisions raise.""" + g = pygraph() + a = g.tensor(dim=[2, 2], uid=2, name="user_uid") # reserve 2 + assert a.uid == 2 and a.uid_assigned + b = g.tensor(dim=[2, 2], name="auto1") # auto: 1 + c = g.tensor(dim=[2, 2], name="auto2") # auto: must skip reserved 2 -> 3 + assert b.uid == 1 + assert c.uid == 3 + d = g.tensor(dim=[2, 2], uid=3, name="steals_from_auto") + assert d.uid == 3 and d.uid_assigned + assert c.uid not in (2, 3) and not c.uid_assigned # renumbered + assert g._tensor_by_uid[c.uid] is c + with pytest.raises(ValueError, match="user-assigned"): + g.tensor(dim=[2, 2], uid=2, name="dup_user") # user-user collides + + def test_matmul(self): + g = pygraph() + A = g.tensor(dim=[8, 64, 128], name="A") + B = g.tensor(dim=[8, 128, 256], name="B") + C = g.matmul(A, B, name="mm1") + + assert len(g.nodes) == 1 + assert g.nodes[0].node_type == NodeType.MATMUL + assert g.nodes[0].name == "mm1" + assert C.is_virtual + + def test_matmul_inputs_outputs(self): + g = pygraph() + A = g.tensor(dim=[8, 64, 128], name="A") + B = g.tensor(dim=[8, 128, 256], name="B") + C = g.matmul(A, B, name="mm1") + + node = g.nodes[0] + assert node.inputs["A"] is A + assert node.inputs["B"] is B + assert node.outputs["C"] is C + assert node.params["padding"] == 0.0 + + def test_find_tensor_by_name(self): + g = pygraph() + t = g.tensor(dim=[8, 64], name="test") + assert g.find_tensor("test") is t + + def test_find_tensor_by_uid(self): + g = pygraph() + t = g.tensor(dim=[8, 64], name="test") + assert g.find_tensor(t.uid) is t + + def test_find_tensor_not_found(self): + g = pygraph() + assert g.find_tensor("nonexistent") is None + + def test_inspect(self): + g = pygraph(io_data_type="HALF") + A = g.tensor(dim=[8, 64], name="A") + B = g.tensor(dim=[64, 32], name="B") + C = g.matmul(A, B, name="mm1") + + info = g.inspect() + assert len(info["nodes"]) == 1 + assert info["nodes"][0]["name"] == "mm1" + assert info["nodes"][0]["type"] == "MATMUL" + assert info["nodes"][0]["params"]["padding"] == 0.0 + assert "A" in info["tensors"] + + def test_auto_naming(self): + g = pygraph() + A = g.tensor(dim=[8, 64], name="A") + B = g.tensor(dim=[64, 32], name="B") + + g.matmul(A, B) + g.matmul(A, B) + + assert g.nodes[0].name == "matmul.0" + assert g.nodes[1].name == "matmul.1" + + def test_validation(self): + g = pygraph() + A = g.tensor(dim=[8, 64], stride=[64, 1], name="A") + B = g.tensor(dim=[64, 32], stride=[32, 1], name="B") + g.matmul(A, B) + g.validate() + + def test_pointwise_add(self): + g = pygraph() + A = g.tensor(dim=[8, 64], name="A") + B = g.tensor(dim=[8, 64], name="B") + C = g.add(A, B) + + assert len(g.nodes) == 1 + assert g.nodes[0].node_type == NodeType.POINTWISE + assert "mode" in g.nodes[0].params + + def test_relu(self): + g = pygraph() + X = g.tensor(dim=[8, 64], name="X") + Y = g.relu(X) + assert g.nodes[0].node_type == NodeType.POINTWISE + + def test_all_pointwise_builders(self): + """Every op in _POINTWISE_TENSOR_ARGS has a builder: positional AND the + classic pybind keyword call styles both produce a first-class node.""" + for op, argnames in pygraph._POINTWISE_TENSOR_ARGS.items(): + for style in ("positional", "keyword"): + g = pygraph() + tensors = [g.tensor(dim=[4, 8], name=f"t{i}") for i in range(len(argnames))] + builder = getattr(g, op) + out = builder(*tensors) if style == "positional" else builder(**dict(zip(argnames, tensors))) + (node,) = g.nodes + assert node.node_type == NodeType.POINTWISE, op + assert node.params["mode"] == op + assert len(node.inputs) == len(argnames), op + assert out.dim == [] or out.dim == [4, 8] # inferred at validate + g.validate() + # classic sequencing lowers (and freezes) at validate: sealed + # dims are tuples — compare by value + assert list(node.outputs["OUT_0"].dim) == [4, 8], op + + def test_all_structured_builders(self): + """Every op in _STRUCTURED_OPS builds a first-class node: named ports + (== C++ kwargs), attrs stored verbatim, declared outputs — via both + keyword and positional-tensor call styles.""" + from cudnn._pygraph import _STRUCTURED_OPS + + for op, spec in _STRUCTURED_OPS.items(): + for style in ("keyword", "positional"): + g = pygraph() + tensors = {port: g.tensor(dim=[4, 8], name=f"{port}_in") for port in spec["inputs"]} + attrs = {ak: "ATTR_SENTINEL" for ak in spec.get("attrs", ())} + lists = {lp: [g.tensor(dim=[4, 8], name=f"{lp}{i}_in") for i in range(2)] for lp in spec.get("list_inputs", ())} + if style == "keyword": + outs = getattr(g, op)(**tensors, **attrs, **lists) + else: + outs = getattr(g, op)(*tensors.values(), **attrs, **lists) + outs = outs if isinstance(outs, (tuple, list)) else (outs,) # classic returns a LIST for multi-output + (node,) = g.nodes + assert node.node_type == spec["node_type"], op + expect_ports = set(spec["inputs"]) | {f"{lp}_{i}" for lp in lists for i in range(2)} + assert set(node.inputs) == expect_ports, op + assert tuple(node.outputs) == spec["outputs"], op + for ak in spec.get("attrs", ()): + assert node.params[ak] == "ATTR_SENTINEL", op + assert len(outs) == len(spec["outputs"]), op + + def test_structured_out_dims(self): + """out_dims sets output dims for shapes cuDNN cannot infer (reduction).""" + g = pygraph() + A = g.tensor(dim=[1, 4, 8], name="A") + R = g.reduction(A, mode="ADD_SENTINEL", out_dims=[1, 4, 1]) + assert R.dim == [1, 4, 1] and R.stride == [4, 1, 1] + + def test_batchnorm_peer_stats_ports(self): + """List inputs (peer_stats) become indexed ports + a count param.""" + g = pygraph() + kwargs = {p: g.tensor(dim=[4, 8], name=p) for p in ("input", "scale", "bias", "epsilon", "momentum", "in_running_mean", "in_running_var")} + ps = [g.tensor(dim=[4, 8], name=f"ps{i}") for i in range(2)] + g.batchnorm(peer_stats=ps, **kwargs) + (node,) = g.nodes + assert node.params["_n_peer_stats"] == 2 + assert "peer_stats_0" in node.inputs and "peer_stats_1" in node.inputs + + def test_pointwise_scalar_attrs(self): + """Ops with scalar attributes store them in params (introspectable).""" + g = pygraph() + X = g.tensor(dim=[4, 8], name="X") + g.relu(X, lower_clip=0.1, upper_clip=6.0) + g.leaky_relu(X, negative_slope=0.01) + g.swish(X, swish_beta=1.5) + g.gen_index(X, axis=1) + r, lr, sw, gi = g.nodes + assert r.params == {"mode": "relu", "lower_clip": 0.1, "upper_clip": 6.0} + assert lr.params == {"mode": "leaky_relu", "negative_slope": 0.01} + assert sw.params == {"mode": "swish", "swish_beta": 1.5} + assert gi.params == {"mode": "gen_index", "axis": 1} + + def test_chaining(self): + g = pygraph() + A = g.tensor(dim=[8, 64, 128], name="A") + B = g.tensor(dim=[8, 128, 256], name="B") + bias = g.tensor(dim=[1, 1, 256], name="bias") + + C = g.matmul(A, B) + D = g.add(C, bias) + E = g.relu(D) + + assert len(g.nodes) == 3 + assert [n.node_type for n in g.nodes] == [NodeType.MATMUL, NodeType.POINTWISE, NodeType.POINTWISE] + + def test_get_node(self): + g = pygraph() + A = g.tensor(dim=[8, 64], name="A") + B = g.tensor(dim=[64, 32], name="B") + g.matmul(A, B, name="mm1") + + node = g.get_node("mm1") + assert node.name == "mm1" + assert g.get_node("nonexistent") is None + + def test_sdpa_inference(self): + """Test SDPA forward inference mode.""" + g = pygraph() + # [B, H, S, D] layout + Q = g.tensor(dim=[2, 8, 128, 64], name="Q") + K = g.tensor(dim=[2, 8, 128, 64], name="K") + V = g.tensor(dim=[2, 8, 128, 64], name="V") + + O, stats = g.sdpa(Q, K, V, is_inference=True, use_causal_mask=True, name="attn") + + assert len(g.nodes) == 1 + assert g.nodes[0].node_type == NodeType.SDPA + assert g.nodes[0].params["is_inference"] is True + assert g.nodes[0].params["use_causal_mask"] is True + assert "O" in g.nodes[0].outputs + assert stats is None # classic API returns [O, None] in inference mode + assert O.dim == [2, 8, 128, 64] # q dims with v's head dim + + def test_sdpa_training(self): + """Test SDPA forward training mode (returns stats).""" + g = pygraph() + Q = g.tensor(dim=[2, 8, 128, 64], name="Q") + K = g.tensor(dim=[2, 8, 128, 64], name="K") + V = g.tensor(dim=[2, 8, 128, 64], name="V") + + O, stats = g.sdpa(Q, K, V, is_inference=False, attn_scale=0.125, name="attn") + + assert len(g.nodes) == 1 + assert g.nodes[0].params["is_inference"] is False + assert g.nodes[0].params["attn_scale"] == 0.125 + assert "O" in g.nodes[0].outputs + assert "Stats" in g.nodes[0].outputs + assert stats.dim == [2, 8, 128, 1] + + +@pytest.mark.L1 +class TestIntegration: + """Integration tests requiring cuDNN.""" + + @pytest.fixture + def cudnn_available(self): + try: + import cudnn + + return cudnn.backend_version() >= 91200 + except Exception: + return False + + def test_build(self, cudnn_available): + if not cudnn_available: + pytest.skip("cuDNN not available") + + import cudnn + + g = pygraph( + io_data_type=cudnn.data_type.HALF, + compute_data_type=cudnn.data_type.FLOAT, + ) + A = g.tensor(dim=[8, 64, 128], name="A") + B = g.tensor(dim=[8, 128, 256], name="B") + C = g.matmul(A, B) + C.set_output(True) + + g.build() + assert g._is_built + assert g.get_workspace_size() >= 0 + + def test_sdpa_build(self, cudnn_available): + """Test building SDPA graph.""" + if not cudnn_available: + pytest.skip("cuDNN not available") + + import cudnn + + g = pygraph( + io_data_type=cudnn.data_type.HALF, + compute_data_type=cudnn.data_type.FLOAT, + ) + # [B, H, S, D] layout + Q = g.tensor(dim=[2, 8, 128, 64], name="Q") + K = g.tensor(dim=[2, 8, 128, 64], name="K") + V = g.tensor(dim=[2, 8, 128, 64], name="V") + + O, _ = g.sdpa(Q, K, V, is_inference=True, use_causal_mask=True, name="attn") + O.set_output(True) + + try: + g.build() + assert g._is_built + except cudnn.cudnnGraphNotSupportedError: + pytest.skip("SDPA not supported on this hardware/configuration") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) + + +class TestReviewSemantics: + """Review items 3 + 4: SDPA port direction; graph-owned identity mutation.""" + + def test_sdpa_output_direction(self): + """dBias & co. are outputs of the node, not inputs (review item 3).""" + g = pygraph() + t = lambda n: g.tensor(dim=[2, 4, 8, 16], name=n) # noqa: E731 + dbias = g.tensor(dim=[1, 4, 8, 8], name="dbias_buf") + g.sdpa_backward(t("q"), t("k"), t("v"), t("o"), t("dO"), t("stats"), dBias=dbias) + (node,) = g.nodes + assert "dBias" in node.outputs and node.outputs["dBias"] is dbias + assert "dBias" not in node.inputs + + def test_tensor_rename_reindexes(self): + g = pygraph() + a = g.tensor(dim=[2, 2], name="old") + a.set_name("new") + assert g.find_tensor("new") is a and g.find_tensor("old") is None + g.tensor(dim=[2, 2], name="other") + a.set_name("other") # classic: labels may collide — becomes ambiguous + with pytest.raises(ValueError, match="ambiguous"): + g.find_tensor("other") + + def test_set_uid_steals_auto_uid_and_rejects_user_dup(self): + """Classic parity: user set_uid wins over an auto-assigned holder (which + is silently renumbered); two USER uids colliding is an error.""" + g = pygraph() + a = torch.randn(2, 2) + A = g.tensor_like(a, name="A") # auto uid 1 + g._data_bindings[A.uid] = a # simulate auto-binding + B = g.tensor(dim=[2, 2], name="B") # auto uid 2 + B.set_uid(A.uid) # user claims A's auto uid + assert B.uid_assigned and g.find_tensor(B.uid) is B + assert A.uid != B.uid and g.find_tensor(A.uid) is A # A renumbered + assert g._data_bindings.get(A.uid) is a # binding followed A + C = g.tensor(dim=[2, 2], name="C") + with pytest.raises(ValueError, match="user-assigned"): + C.set_uid(B.uid) + + def test_tensor_dict_key_stable_across_mutation(self): + """Identity-based hashing: a Tensor used as a dict key survives + uid/name mutation (review item 4).""" + g = pygraph() + A = g.tensor(dim=[2, 2], name="A") + d = {A: "x"} + A.set_name("renamed") + A.set_uid(1000) + assert d[A] == "x" + + def test_identity_mutation_frozen_after_planning(self): + from cudnn.engines import BaseEngine, PYTHON_ENGINE_ID_BASE + + class Dummy(BaseEngine): + engine_id = PYTHON_ENGINE_ID_BASE + 90 + + def execute(self, graph, tensor_data, ctx=None): + pass + + g = pygraph() + g.register_backend(Dummy()) # keeps planning python-side (no C++ needed) + A = g.tensor(dim=[1, 2, 2], name="A") + g.matmul(A, g.tensor(dim=[1, 2, 2], name="B")) + g.create_execution_plans() + with pytest.raises(RuntimeError, match="frozen"): + A.set_uid(500) + + def test_mxfp8_dsink_is_output(self): + """Follow-up item 4: mxfp8_backward dSink_token is an output port.""" + g = pygraph() + t = lambda n: g.tensor(dim=[2, 4, 8, 16], name=n) # noqa: E731 + kw = {p: t(p) for p in ("q", "q_T", "k", "k_T", "v", "o_f16", "dO_f16", "dO", "dO_T", "stats")} + ds = g.tensor(dim=[1, 4, 1, 1], name="dsink_buf") + g.sdpa_mxfp8_backward(dSink_token=ds, **kw) + (node,) = g.nodes + assert "dSink_token" in node.outputs and "dSink_token" not in node.inputs + + def test_semantic_setters_frozen_after_planning(self): + from cudnn.engines import BaseEngine, PYTHON_ENGINE_ID_BASE + + class Dummy(BaseEngine): + engine_id = PYTHON_ENGINE_ID_BASE + 91 + + def execute(self, graph, tensor_data, ctx=None): + pass + + g = pygraph(backends=[Dummy()]) + A = g.tensor(dim=[1, 2, 2], name="A") + g.matmul(A, g.tensor(dim=[1, 2, 2], name="B")) + g.create_execution_plans() + for mutate in (lambda: A.set_dim([4, 4]), lambda: A.set_data_type("HALF"), lambda: A.set_output(True), lambda: A.set_stride([4, 1])): + with pytest.raises(RuntimeError, match="frozen"): + mutate() + + def test_freeze_covers_public_surface(self): + """Review round 5: the freeze must close EVERY public mutation path, + not only the fluent API — attribute writes, live containers, in-place + list edits, node params, and graph context.""" + from cudnn.engines import BaseEngine, PYTHON_ENGINE_ID_BASE + + class Dummy(BaseEngine): + engine_id = PYTHON_ENGINE_ID_BASE + 92 + + def execute(self, graph, tensor_data, ctx=None): + pass + + g = pygraph(backends=[Dummy()]) + A = g.tensor(dim=[1, 2, 2], name="A") + C = g.matmul(A, g.tensor(dim=[1, 2, 2], name="B")) + g.create_execution_plans() + node = g.nodes[0] + + with pytest.raises(RuntimeError, match="frozen"): + A.dim = [9, 9] # direct attribute write + with pytest.raises(TypeError): + A.dim[:] = [9] # sealed to a tuple: no in-place edits + with pytest.raises(TypeError): + node.params["padding"] = 123 # MappingProxy + with pytest.raises(TypeError): + node.inputs["A"] = C # MappingProxy + with pytest.raises(RuntimeError, match="frozen"): + node.inputs = {} # attribute write on the node + with pytest.raises(RuntimeError, match="frozen"): + g.context.compute_data_type = "HALF" # graph context + # live-container laundering: the public views are copies + g.nodes.clear() + g.tensors.clear() + assert len(g.nodes) == 1 and len(g.tensors) == 3 + # the inspection surface stays readable for engines + assert list(node.inputs) == ["A", "B"] and list(C.dim) == [1, 2, 2] + + def test_mutation_after_validate_revalidates(self): + """Review round 5: python-engine graphs stay mutable until planning — + but a mutation after validate() must invalidate _is_validated so stale + inference never reaches planning.""" + from cudnn.engines import BaseEngine, PYTHON_ENGINE_ID_BASE + + class Dummy(BaseEngine): + engine_id = PYTHON_ENGINE_ID_BASE + 93 + + def execute(self, graph, tensor_data, ctx=None): + pass + + g = pygraph(backends=[Dummy()]) + A = g.tensor(dim=[1, 2, 2], name="A") + g.matmul(A, g.tensor(dim=[1, 2, 2], name="B")) + g.validate() + assert g._is_validated and not g._frozen # mutable until planning + A.set_data_type("HALF") # allowed — and must force re-validation + assert not g._is_validated + g.create_execution_plans() + assert g._frozen + + def test_tensor_scalar_is_graph_owned(self): + g = pygraph() + s = g.tensor_scalar(1.5, scalar_type="FLOAT_SENTINEL") + s.set_name("renamed_scalar") + assert g.find_tensor("renamed_scalar") is s + + def test_duplicate_names_are_classic_labels(self): + """Classic parity: tensor names are debug labels — duplicates are legal + (pycudnnTest builds two tensors both named 'weight'). uid is the + identity; name-keyed lookups on an ambiguous label raise instead of + guessing, unique labels keep working.""" + g = pygraph() + a = g.tensor(dim=[2, 2], name="X") + b = g.tensor(dim=[2, 2], name="X") # legal, like classic + assert a.name == b.name == "X" and a.uid != b.uid + with pytest.raises(ValueError, match="ambiguous"): + g.find_tensor("X") + assert g.find_tensor(a.uid) is a and g.find_tensor(b.uid) is b + u = g.tensor(dim=[2, 2], name="unique") + assert g.find_tensor("unique") is u + # renaming ONTO an existing label is equally legal — and makes it ambiguous + u.set_name("X") + with pytest.raises(ValueError, match="ambiguous"): + g.find_tensor("X") diff --git a/test/python/test_native_backend_lowering.py b/test/python/test_native_backend_lowering.py new file mode 100644 index 000000000..594dee190 --- /dev/null +++ b/test/python/test_native_backend_lowering.py @@ -0,0 +1,519 @@ +"""GPU parity: pygraph builds natively, lowers to cuDNN, executes correctly. + +Covers the native -> _lower_to_cpp -> cuDNN execute path (uid propagation, handle +threading, pointwise dispatch). Skipped without a GPU / cuDNN. +""" + +import pytest + +torch = pytest.importorskip("torch") +if not torch.cuda.is_available(): + pytest.skip("needs a CUDA GPU", allow_module_level=True) + +import cudnn +from cudnn._pygraph import pygraph + +pytestmark = pytest.mark.L0 + +M, K, N = 64, 32, 48 + + +def _handle(): + return cudnn.create_handle() + + +def _assert_ran_on_backend(g): + """Dispatch-level proof the execution took the cuDNN backend plan path: + the selected routed plan is the backend entry (no python engine), the graph + was really lowered, and backend plans were created and built. Kernel + identity below the backend API is deliberately not asserted (kernel names + are backend-internal and version-dependent).""" + assert g.selected_engine is None + assert g._lowered_graph is not None + assert g._cpp_plans_created and g._is_built + + +def test_native_matmul_lowers_to_backend(): + h = _handle() + a = torch.randn(1, M, K, device="cuda", dtype=torch.float16) + b = torch.randn(1, K, N, device="cuda", dtype=torch.float16) + c = torch.empty(1, M, N, device="cuda", dtype=torch.float16) + + g = pygraph(handle=h, io_data_type=cudnn.data_type.HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + A = g.tensor(dim=[1, M, K], stride=[M * K, K, 1], data_type=cudnn.data_type.HALF) + B = g.tensor(dim=[1, K, N], stride=[K * N, N, 1], data_type=cudnn.data_type.HALF) + C = g.matmul(A, B) + C.set_output(True).set_data_type(cudnn.data_type.HALF) + + g.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + g.execute({A: a, B: b, C: c}, ws, handle=h) + torch.cuda.synchronize() + _assert_ran_on_backend(g) + + torch.testing.assert_close(c.float(), a.float() @ b.float(), atol=2e-2, rtol=2e-2) + + +def test_native_matmul_bias_relu_lowers_to_backend(): + h = _handle() + a = torch.randn(1, M, K, device="cuda", dtype=torch.float16) + b = torch.randn(1, K, N, device="cuda", dtype=torch.float16) + bias = torch.randn(1, M, N, device="cuda", dtype=torch.float16) + c = torch.empty(1, M, N, device="cuda", dtype=torch.float16) + + g = pygraph(handle=h, io_data_type=cudnn.data_type.HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + A = g.tensor(dim=[1, M, K], stride=[M * K, K, 1], data_type=cudnn.data_type.HALF) + B = g.tensor(dim=[1, K, N], stride=[K * N, N, 1], data_type=cudnn.data_type.HALF) + Bi = g.tensor(dim=[1, M, N], stride=[M * N, N, 1], data_type=cudnn.data_type.HALF) + Y = g.relu(g.bias(g.matmul(A, B), Bi)) + Y.set_output(True).set_data_type(cudnn.data_type.HALF) + + g.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + g.execute({A: a, B: b, Bi: bias, Y: c}, ws, handle=h) + torch.cuda.synchronize() + _assert_ran_on_backend(g) + + torch.testing.assert_close(c.float(), torch.relu(a.float() @ b.float() + bias.float()), atol=2e-2, rtol=2e-2) + + +def test_native_matmul_reduction_lowers_to_backend(): + """matmul -> reduction(ADD) over N; cuDNN needs explicit reduced output dims.""" + h = _handle() + a = torch.randn(1, M, K, device="cuda", dtype=torch.float16) + b = torch.randn(1, K, N, device="cuda", dtype=torch.float16) + r = torch.empty(1, M, 1, device="cuda", dtype=torch.float32) + + g = pygraph(handle=h, io_data_type=cudnn.data_type.HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + A = g.tensor(dim=[1, M, K], stride=[M * K, K, 1], data_type=cudnn.data_type.HALF) + B = g.tensor(dim=[1, K, N], stride=[K * N, N, 1], data_type=cudnn.data_type.HALF) + R = g.reduction(g.matmul(A, B), mode=cudnn.reduction_mode.ADD, out_dims=[1, M, 1]) + R.set_output(True).set_data_type(cudnn.data_type.FLOAT) + + g.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + g.execute({A: a, B: b, R: r}, ws, handle=h) + torch.cuda.synchronize() + _assert_ran_on_backend(g) + + torch.testing.assert_close(r, (a.float() @ b.float()).sum(dim=2, keepdim=True), atol=5e-2, rtol=5e-2) + + +def test_native_block_scale_nvfp4_lowers_to_backend(): + """block_scale_dequantize(A)@block_scale_dequantize(B), nvfp4 -> cuDNN (SM100).""" + if not hasattr(torch, "float4_e2m1fn_x2"): + pytest.skip("torch lacks float4_e2m1fn_x2") + if torch.cuda.get_device_properties(0).major < 10: + pytest.skip("block-scale MMA needs SM100+") + + h = _handle() + b, Mb, Nb, Kb, BS = 1, 128, 128, 64, 16 + A = torch.randint(0, 256, (b, Mb, Kb // 2), dtype=torch.uint8, device="cuda").view(torch.float4_e2m1fn_x2) + B = torch.randint(0, 256, (b, Kb, Nb // 2), dtype=torch.uint8, device="cuda").view(torch.float4_e2m1fn_x2) + k_scale = ((Kb + BS - 1) // BS + 3) // 4 * 4 + A_ds = torch.full((b, 128, k_scale), 1.0, dtype=torch.float8_e4m3fn, device="cuda") + B_ds = torch.full((b, k_scale, 128), 1.0, dtype=torch.float8_e4m3fn, device="cuda") + C = torch.empty((b, Mb, Nb), dtype=torch.bfloat16, device="cuda") + + g = pygraph(handle=h, compute_data_type=cudnn.data_type.FLOAT) + At = g.tensor(dim=[b, Mb, Kb], stride=[Mb * Kb, Kb, 1], data_type=cudnn.data_type.FP4_E2M1) + Bt = g.tensor(dim=[b, Kb, Nb], stride=[Nb * Kb, 1, Kb], data_type=cudnn.data_type.FP4_E2M1) + Ad = g.tensor( + dim=[b, 128, k_scale], stride=[128 * k_scale, k_scale, 1], data_type=cudnn.data_type.FP8_E4M3, reordering_type=cudnn.tensor_reordering.F8_128x4 + ) + Bd = g.tensor( + dim=[b, k_scale, 128], stride=[k_scale * 128, 1, k_scale], data_type=cudnn.data_type.FP8_E4M3, reordering_type=cudnn.tensor_reordering.F8_128x4 + ) + Cc = g.matmul( + g.block_scale_dequantize(At, Ad, block_size=[1, BS]), g.block_scale_dequantize(Bt, Bd, block_size=[BS, 1]), compute_data_type=cudnn.data_type.FLOAT + ) + Cc.set_output(True).set_data_type(cudnn.data_type.BFLOAT16) + + g.build([cudnn.heur_mode.A, cudnn.heur_mode.B]) + ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + g.execute({At: A, Bt: B, Ad: A_ds, Bd: B_ds, Cc: C}, ws, handle=h) + torch.cuda.synchronize() + _assert_ran_on_backend(g) # builds + executes without error (parity harness = repo's fp4 test) + + +def test_native_moe_grouped_matmul_lowers_to_backend(): + """moe_grouped_matmul (mode=NONE) built natively -> cuDNN, parity vs a + self-contained per-expert reference.""" + if cudnn.backend_version() < 91500: + pytest.skip("moe_grouped_matmul requires cuDNN 9.15+") + h = _handle() + E, T, Wt, Hd = 8, 256, 64, 128 + fto = [i * (T // E) for i in range(E)] # one contiguous token chunk per expert + + g = pygraph(handle=h, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + tok = g.tensor(dim=[1, T, Hd], stride=[T * Hd, Hd, 1], data_type=cudnn.data_type.BFLOAT16) + wt = g.tensor(dim=[E, Hd, Wt], stride=[Hd * Wt, 1, Hd], data_type=cudnn.data_type.BFLOAT16) + off = g.tensor(dim=[E, 1, 1], stride=[1, 1, 1], data_type=cudnn.data_type.INT32) + out = g.moe_grouped_matmul(tok, wt, off, mode=cudnn.moe_grouped_matmul_mode.NONE, compute_data_type=cudnn.data_type.FLOAT) + out.set_data_type(cudnn.data_type.BFLOAT16).set_output(True) + + g.build([cudnn.heur_mode.A]) + tok_d = torch.randn(T * Hd, dtype=torch.bfloat16, device="cuda") + wt_d = torch.randn(E * Hd * Wt, dtype=torch.bfloat16, device="cuda") + off_d = torch.tensor(fto, dtype=torch.int32, device="cuda") + out_d = torch.empty(T * Wt, dtype=torch.bfloat16, device="cuda") + ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") + g.execute({tok: tok_d, wt: wt_d, off: off_d, out: out_d}, ws, handle=h) + torch.cuda.synchronize() + _assert_ran_on_backend(g) + + # reference: per-expert token-chunk @ weight[e] (weights stored H-contiguous) + token = tok_d.view(T, Hd).float() + weight = torch.as_strided(wt_d.float(), (E, Hd, Wt), (Hd * Wt, 1, Hd)) + ref = torch.empty(T, Wt) + bounds = fto + [T] + for e in range(E): + s, en = bounds[e], bounds[e + 1] + if en > s: + ref[s:en] = token[s:en] @ weight[e] + torch.testing.assert_close(out_d.view(T, Wt).float(), ref.cuda(), rtol=5e-2, atol=5e-2) + + +def test_native_sdpa_fwd_lowers_to_backend(): + """sdpa (captured-op family) -> cuDNN execution parity vs torch SDPA.""" + h = _handle() + B, Hh, S, D = 2, 4, 128, 64 + q = torch.randn(B, Hh, S, D, device="cuda", dtype=torch.float16) + k = torch.randn(B, Hh, S, D, device="cuda", dtype=torch.float16) + v = torch.randn(B, Hh, S, D, device="cuda", dtype=torch.float16) + o = torch.empty(B, Hh, S, D, device="cuda", dtype=torch.float16) + ref = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True) + + g = pygraph(handle=h, io_data_type=cudnn.data_type.HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + Q = g.tensor(dim=[B, Hh, S, D], stride=list(q.stride()), data_type=cudnn.data_type.HALF) + K = g.tensor(dim=[B, Hh, S, D], stride=list(k.stride()), data_type=cudnn.data_type.HALF) + V = g.tensor(dim=[B, Hh, S, D], stride=list(v.stride()), data_type=cudnn.data_type.HALF) + O, stats = g.sdpa(Q, K, V, is_inference=True, use_causal_mask=True, attn_scale=1.0 / (D**0.5)) + assert stats is None and O.dim == [B, Hh, S, D] + O.set_output(True).set_data_type(cudnn.data_type.HALF) + + g.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + g.execute({Q: q, K: k, V: v, O: o}, ws, handle=h) + torch.cuda.synchronize() + _assert_ran_on_backend(g) + + torch.testing.assert_close(o, ref, atol=5e-2, rtol=5e-2) + + +def test_native_conv_fprop_lowers_to_backend(): + """conv_fprop (structured-table op) -> cuDNN parity vs torch conv2d (NHWC).""" + h = _handle() + x = torch.randn(4, 16, 32, 32, device="cuda", dtype=torch.float16).to(memory_format=torch.channels_last) + w = torch.randn(32, 16, 3, 3, device="cuda", dtype=torch.float16).to(memory_format=torch.channels_last) + ref = torch.nn.functional.conv2d(x, w, padding=[1, 1], stride=[1, 1], dilation=[1, 1]) + y = torch.empty_like(ref).to(memory_format=torch.channels_last) + + g = pygraph(handle=h, io_data_type=cudnn.data_type.HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + X = g.tensor(dim=list(x.shape), stride=list(x.stride()), data_type=cudnn.data_type.HALF) + W = g.tensor(dim=list(w.shape), stride=list(w.stride()), data_type=cudnn.data_type.HALF) + Y = g.conv_fprop(image=X, weight=W, padding=[1, 1], stride=[1, 1], dilation=[1, 1]) + assert Y.dim == list(ref.shape) # table shape inference + Y.set_output(True).set_data_type(cudnn.data_type.HALF) + + g.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + g.execute({X: x, W: w, Y: y}, ws, handle=h) + torch.cuda.synchronize() + _assert_ran_on_backend(g) + + torch.testing.assert_close(y, ref, atol=5e-2, rtol=5e-2) + + +def test_native_layernorm_fwd_bwd_lowers_to_backend(): + """layernorm fwd (3 outputs) + layernorm_backward (3 outputs) through the + generic structured-op lowering, parity vs torch autograd. + + Uses the cuDNN-supported LN config ([N, C, 1, 1] channels_last, i.e. LN over + the embedding dim) — same as the classic test_layernorm.py.""" + h = _handle() + Nb, C = 64, 128 + eps = 1e-3 + + def cl(t): + return t.to(memory_format=torch.channels_last) + + x = cl(torch.randn(Nb, C, 1, 1, device="cuda", dtype=torch.float16)).requires_grad_() + scale = cl(torch.randn(1, C, 1, 1, device="cuda", dtype=torch.float16)).requires_grad_() + bias = cl(torch.randn(1, C, 1, 1, device="cuda", dtype=torch.float16)).requires_grad_() + eps_cpu = torch.full((1, 1, 1, 1), eps, dtype=torch.float32) + + # torch reference (normalize over all non-batch dims) + xf = x.float() + mean_ref = xf.mean(dim=(1, 2, 3), keepdim=True) + inv_ref = torch.rsqrt(xf.var(dim=(1, 2, 3), keepdim=True, unbiased=False) + eps) + Y_ref = (xf - mean_ref) * inv_ref * scale.float() + bias.float() + grad = torch.randn_like(Y_ref) + Y_ref.backward(grad) + + cl_stride = [C, 1, C, C] # channels_last for [*, C, 1, 1] + + # ---- forward ---- + g = pygraph(handle=h, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + X = g.tensor(dim=[Nb, C, 1, 1], stride=cl_stride, data_type=cudnn.data_type.HALF) + S = g.tensor(dim=[1, C, 1, 1], stride=cl_stride, data_type=cudnn.data_type.HALF) + Bi = g.tensor(dim=[1, C, 1, 1], stride=cl_stride, data_type=cudnn.data_type.HALF) + E = g.tensor(dim=[1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.FLOAT, is_pass_by_value=True) + Y, mean, iv = g.layernorm(norm_forward_phase=cudnn.norm_forward_phase.TRAINING, input=X, scale=S, bias=Bi, epsilon=E) + Y.set_output(True).set_data_type(cudnn.data_type.HALF) + mean.set_output(True).set_data_type(cudnn.data_type.FLOAT) + iv.set_output(True).set_data_type(cudnn.data_type.FLOAT) + g.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + Yb = cl(torch.empty(Nb, C, 1, 1, device="cuda", dtype=torch.float16)) + mb = torch.empty(Nb, 1, 1, 1, device="cuda", dtype=torch.float32) + ivb = torch.empty(Nb, 1, 1, 1, device="cuda", dtype=torch.float32) + ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + g.execute({X: x.detach(), S: scale.detach(), Bi: bias.detach(), E: eps_cpu, Y: Yb, mean: mb, iv: ivb}, ws, handle=h) + torch.cuda.synchronize() + _assert_ran_on_backend(g) + torch.testing.assert_close(Yb.float(), Y_ref, atol=3e-2, rtol=3e-2) + torch.testing.assert_close(mb, mean_ref, atol=5e-3, rtol=5e-3) + torch.testing.assert_close(ivb, inv_ref, atol=5e-3, rtol=5e-3) + + # ---- backward ---- + g2 = pygraph(handle=h, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + DY = g2.tensor(dim=[Nb, C, 1, 1], stride=cl_stride, data_type=cudnn.data_type.HALF) + X2 = g2.tensor(dim=[Nb, C, 1, 1], stride=cl_stride, data_type=cudnn.data_type.HALF) + S2 = g2.tensor(dim=[1, C, 1, 1], stride=cl_stride, data_type=cudnn.data_type.HALF) + M2 = g2.tensor(dim=[Nb, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.FLOAT) + IV2 = g2.tensor(dim=[Nb, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.FLOAT) + DX, DS, DB = g2.layernorm_backward(grad=DY, input=X2, scale=S2, mean=M2, inv_variance=IV2) + for t in (DX, DS, DB): + t.set_output(True).set_data_type(cudnn.data_type.HALF) + g2.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + dxb = cl(torch.empty(Nb, C, 1, 1, device="cuda", dtype=torch.float16)) + dsb = cl(torch.empty(1, C, 1, 1, device="cuda", dtype=torch.float16)) + dbb = cl(torch.empty(1, C, 1, 1, device="cuda", dtype=torch.float16)) + ws2 = torch.empty(max(g2.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + g2.execute({DY: cl(grad.half()), X2: x.detach(), S2: scale.detach(), M2: mb, IV2: ivb, DX: dxb, DS: dsb, DB: dbb}, ws2, handle=h) + torch.cuda.synchronize() + _assert_ran_on_backend(g2) + torch.testing.assert_close(dxb.float(), x.grad.float(), atol=5e-2, rtol=5e-2) + torch.testing.assert_close(dsb.float(), scale.grad.float(), atol=5e-2, rtol=5e-2) + torch.testing.assert_close(dbb.float(), bias.grad.float(), atol=5e-2, rtol=5e-2) + + +def test_native_pointwise_batch_lowers_to_backend(): + """Generated pointwise builders through real cuDNN: sqrt(abs(A@B)) clamped + via binary max/min (keyword call style, input0/input1).""" + h = _handle() + a = torch.randn(1, M, K, device="cuda", dtype=torch.float16) + b = torch.randn(1, K, N, device="cuda", dtype=torch.float16) + lo = torch.full((1, 1, 1), 0.5, device="cuda", dtype=torch.float32) + hi = torch.full((1, 1, 1), 2.0, device="cuda", dtype=torch.float32) + c = torch.empty(1, M, N, device="cuda", dtype=torch.float16) + + g = pygraph(handle=h, io_data_type=cudnn.data_type.HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + A = g.tensor(dim=[1, M, K], stride=[M * K, K, 1], data_type=cudnn.data_type.HALF) + B = g.tensor(dim=[1, K, N], stride=[K * N, N, 1], data_type=cudnn.data_type.HALF) + Lo = g.tensor(dim=[1, 1, 1], stride=[1, 1, 1], data_type=cudnn.data_type.FLOAT) + Hi = g.tensor(dim=[1, 1, 1], stride=[1, 1, 1], data_type=cudnn.data_type.FLOAT) + Y = g.min(input0=g.max(input0=g.sqrt(g.abs(g.matmul(A, B))), input1=Lo), input1=Hi) + Y.set_output(True).set_data_type(cudnn.data_type.HALF) + + g.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + g.execute({A: a, B: b, Lo: lo, Hi: hi, Y: c}, ws, handle=h) + torch.cuda.synchronize() + _assert_ran_on_backend(g) + + ref = (a.float() @ b.float()).abs().sqrt().clamp(0.5, 2.0) + torch.testing.assert_close(c.float(), ref, atol=2e-2, rtol=2e-2) + + +def test_native_rmsnorm_lowers_to_backend(): + """rmsnorm (multi-output: Y + inv_var, pass-by-value epsilon) -> cuDNN parity. + + Regression cover for uid ownership: the Python IR assigns every uid eagerly + and lowering pushes them all explicitly (set_uid on op outputs), so the C++ + FE's build-time auto-assignment never runs. Without this, multi-output ops + get C++ uids in FE enumeration order (inv_var before Y here) != IR order — + keying the variant pack by IR uids then bound Y's buffer to inv_var + (heap corruption / NaN).""" + h = _handle() + Nb, C, Hh, W = 4, 8, 4, 4 + eps = 1e-3 + x = torch.randn(Nb, C, Hh, W, device="cuda", dtype=torch.float16) + scale = torch.randn(1, C, Hh, W, device="cuda", dtype=torch.float16) + bias = torch.randn(1, C, Hh, W, device="cuda", dtype=torch.float16) + eps_cpu = torch.full((1, 1, 1, 1), eps, dtype=torch.float32) + Yb = torch.empty_like(x) + ivb = torch.empty(Nb, 1, 1, 1, device="cuda", dtype=torch.float32) + + g = pygraph(handle=h, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + X = g.tensor(dim=[Nb, C, Hh, W], stride=[C * Hh * W, Hh * W, W, 1], data_type=cudnn.data_type.HALF) + S = g.tensor(dim=[1, C, Hh, W], stride=[C * Hh * W, Hh * W, W, 1], data_type=cudnn.data_type.HALF) + Bi = g.tensor(dim=[1, C, Hh, W], stride=[C * Hh * W, Hh * W, W, 1], data_type=cudnn.data_type.HALF) + E = g.tensor(dim=[1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.FLOAT, is_pass_by_value=True) + Y, iv = g.rmsnorm(input=X, scale=S, epsilon=E, bias=Bi, norm_forward_phase=cudnn.norm_forward_phase.TRAINING) + Y.set_output(True).set_data_type(cudnn.data_type.HALF) + iv.set_output(True).set_data_type(cudnn.data_type.FLOAT) + + # first-class introspection: named ports + params + node = g.nodes[0] + assert node.node_type.name == "RMSNORM" + assert set(node.inputs) == {"input", "scale", "epsilon", "bias"} + assert set(node.outputs) == {"Y", "inv_var"} + + g.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + g.execute({X: x, S: scale, Bi: bias, E: eps_cpu, Y: Yb, iv: ivb}, ws, handle=h) + torch.cuda.synchronize() + _assert_ran_on_backend(g) + + xf = x.float() + ivref = torch.rsqrt(xf.pow(2).mean(dim=(1, 2, 3), keepdim=True) + eps) + Yref = scale.float() * (xf * ivref) + bias.float() + torch.testing.assert_close(Yb.float(), Yref, atol=3e-2, rtol=3e-2) + torch.testing.assert_close(ivb, ivref, atol=5e-3, rtol=5e-3) + + +def test_mixed_router_backend_slot_executes(): + """Review round 4: the backend entry of a MIXED router is selectable and + actually executes through the backend (lowering triggered), with routed + indices stable across that lowering; the pinned python plan still runs + afterwards with its own knobs.""" + from cudnn.engines import BaseEngine, PlanConfig, Router + from cudnn.engines.engine_ids import BACKEND_HEURISTIC_ENGINE_ID, PYTHON_ENGINE_ID_BASE + + ran = [] + + class PyMatmul(BaseEngine): + name = "py_matmul" + engine_id = PYTHON_ENGINE_ID_BASE + 90 + + def execute(self, graph, tensor_data, ctx=None): + node = graph.nodes[0] + a = tensor_data[node.inputs["A"].uid] + b = tensor_data[node.inputs["B"].uid] + c = tensor_data[node.outputs["C"].uid] + c.copy_((a.float() @ b.float()).to(c.dtype)) + ran.append("python") + + class CudnnFirst(Router): + def plan(self, graph, backends): + return [PlanConfig(BACKEND_HEURISTIC_ENGINE_ID), PlanConfig(backends[0].engine_id)] + + h = _handle() + a = torch.randn(1, M, K, device="cuda", dtype=torch.float16) + b = torch.randn(1, K, N, device="cuda", dtype=torch.float16) + c = torch.empty(1, M, N, device="cuda", dtype=torch.float16) + ref = (a.float() @ b.float()).half() + + g = pygraph( + handle=h, io_data_type=cudnn.data_type.HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT, router=CudnnFirst() + ) + g.register_backend(PyMatmul()) + A = g.tensor(dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(dim=[1, K, N], stride=[K * N, N, 1]) + C = g.matmul(A, B) + C.set_output(True).set_data_type(cudnn.data_type.HALF) + + g.create_execution_plans() + assert [p.engine_id for p in g.plans] == [BACKEND_HEURISTIC_ENGINE_ID, PYTHON_ENGINE_ID_BASE + 90] + + # slot 0 = cuDNN: this build/execute lowers and runs the real backend + assert g.selected_engine is None + g.build() + assert g._lowered_graph is not None # lowering really happened + ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + g.execute({A: a, B: b, C: c}, ws, handle=h) + torch.cuda.synchronize() + _assert_ran_on_backend(g) + torch.testing.assert_close(c.float(), ref.float(), atol=2e-2, rtol=2e-2) + assert ran == [] # the python engine did NOT run + + # backend count is the classic passthrough space; routed indices unmoved + assert g.get_execution_plan_count() >= 1 + assert [p.engine_id for p in g.plans] == [BACKEND_HEURISTIC_ENGINE_ID, PYTHON_ENGINE_ID_BASE + 90] + + # slot 1 = the python plan, still selectable AFTER backend lowering + c.zero_() + g.select_plan(1) + assert g.selected_engine.name == "py_matmul" + g.execute({A: a, B: b, C: c}, ws, handle=h) + torch.cuda.synchronize() + assert g.selected_engine is not None # this slot is the python engine + assert ran == ["python"] + torch.testing.assert_close(c.float(), ref.float(), atol=2e-2, rtol=2e-2) + + +def test_planning_one_shot_backend_only(): + """Review round 4: one-shot planning also covers the pure-cuDNN graph (no + python engines registered) — a second create_execution_plans() raises.""" + fn = pygraph.create_execution_plans + if "pygraph" not in getattr(fn, "__qualname__", ""): + pytest.skip(f"cudnn.pygraph.create_execution_plans is monkey-patched ({getattr(fn, '__qualname__', '?')}); the wrapper swallows the one-shot error") + h = _handle() + g = pygraph(handle=h, io_data_type=cudnn.data_type.HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + A = g.tensor(dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(dim=[1, K, N], stride=[K * N, N, 1]) + C = g.matmul(A, B) + C.set_output(True).set_data_type(cudnn.data_type.HALF) + + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + with pytest.raises(RuntimeError, match="one-shot"): + g.create_execution_plans([cudnn.heur_mode.A]) + # the first plan set is intact and usable + g.check_support() + g.build_plans() + a = torch.randn(1, M, K, device="cuda", dtype=torch.float16) + b = torch.randn(1, K, N, device="cuda", dtype=torch.float16) + c = torch.empty(1, M, N, device="cuda", dtype=torch.float16) + ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + g.execute({A: a, B: b, C: c}, ws, handle=h) + torch.cuda.synchronize() + _assert_ran_on_backend(g) + torch.testing.assert_close(c.float(), (a.float() @ b.float()), atol=2e-2, rtol=2e-2) + + +def test_output_layout_contract(): + """Review round 5: USER-assigned output dim/stride must reach the lowered + C++ tensor; IR-INFERRED strides must NOT be pushed — the backend keeps its + classic per-op layout inference (channels-last conv) when the user did not + pin one. Checked on the lowered graph JSON and by execution.""" + import json + + # (a) explicit matmul output stride is honored end to end + h = _handle() + g = pygraph(handle=h, io_data_type=cudnn.data_type.HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + A = g.tensor(dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(dim=[1, K, N], stride=[K * N, N, 1]) + C = g.matmul(A, B) + C.set_output(True).set_data_type(cudnn.data_type.HALF).set_dim([1, M, N]).set_stride([M * N, 1, M]) # column-major + g.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + + lowered = json.loads(str(g._lowered_graph)) + (c_entry,) = [t for t in lowered["tensors"].values() if t["uid"] == C.uid] + assert c_entry["stride"] == [M * N, 1, M] # user layout pushed verbatim + + a = torch.randn(1, M, K, device="cuda", dtype=torch.float16) + b = torch.randn(1, K, N, device="cuda", dtype=torch.float16) + c = torch.empty(1, N, M, device="cuda", dtype=torch.float16).permute(0, 2, 1) # column-major buffer + ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + g.execute({A: a, B: b, C: c}, ws, handle=h) + torch.cuda.synchronize() + _assert_ran_on_backend(g) + torch.testing.assert_close(c.float(), (a.float() @ b.float()), atol=2e-2, rtol=2e-2) + + # (b) inferred conv output keeps the backend's channels-last inference + # (the IR's provisional row-major stride must NOT leak into C++) + h2 = _handle() + Nn, Cc, Hh, Ww, Kk = 4, 32, 16, 16, 16 + g2 = pygraph(handle=h2, io_data_type=cudnn.data_type.HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + X = g2.tensor(dim=[Nn, Cc, Hh, Ww], stride=[Cc * Hh * Ww, 1, Cc * Ww, Cc]) # NHWC + W = g2.tensor(dim=[Kk, Cc, 3, 3], stride=[Cc * 9, 1, Cc * 3, Cc]) + Y = g2.conv_fprop(X, W, padding=[1, 1], stride=[1, 1], dilation=[1, 1]) + Y.set_output(True).set_data_type(cudnn.data_type.HALF) + g2.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + + lowered2 = json.loads(str(g2._lowered_graph)) + (y_entry,) = [t for t in lowered2["tensors"].values() if t["uid"] == Y.uid] + assert y_entry["stride"][1] == 1, y_entry["stride"] # channels-last kept, not row-major