From 3e726f0dd5931c2f2944fdc474fc6c36503d95f1 Mon Sep 17 00:00:00 2001 From: ysun67 Date: Tue, 19 May 2026 21:37:03 -0400 Subject: [PATCH] PR-P0: pragma plugin framework primitives MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Foundation for Phase P (the pragma-plugin contract, spec § 3.2.1.3 / § 6.0.0). Adds the seven new primitives that unblock the per-pragma migration PRs (PR-P1..P4) and the consequence-cleanup (PR-P5): 1. `PragmaPlugin` schema + `@pragma_plugin` decorator factory (NEW `core/pragma_plugin.py`) 2. `Services` typed-key dict (NEW `core/services.py`) — Risk 1 fix 3. `Attribute` + `AttributeDict` typed-key dict (NEW `core/attribute.py`) — Risks 2 + 5 fix 4. `Op.attributes` field (extends `core/ops.py`) — Risks 2 + 5 applied. `kw_only=True, compare=False, hash=False, repr=False` keeps every existing Op subclass byte-equivalent: positional construction stable, equality / hashing / repr unchanged. 5. Per-Compiler registries: `compiler.{ops, attributes, passes, lowerings, renders, services, pragma_plugins}` (extends `core/dialect.py`) — Risk 4 fix 6. `Compiler.register_pragma_plugin(plugin)` — atomic conflict detection (op-name, attribute-name, render-double-registration, pass-cycle, missing required service) with try-everything-then- commit semantics 7. `Pass.{produces_ops, consumes_ops, preserves}` typed deps + `topo_sort_passes` LLVM-style scheduler + `PassCycleError` (extends `core/passes.py`) 8. `_compiler_registration_scope(compiler)` thread-local sets the "current compiler"; `@pragma_handler` stages into that compiler's per-instance handler list when active, falls back to the module- global registry otherwise (back-compat shim, deleted in PR-P5) 9. `MaterializePragmaPass` back-compat shim Pass that runs the staged pragma handlers in `compiler.run(prog)` Naming note: `Pass.produces`/`consumes` were already str-typed (dialect names) before PR-P0; the new op-type-typed deps are introduced under suffixed names (`produces_ops`, `consumes_ops`, `preserves`) so the two ordering surfaces coexist cleanly during the per-pragma migration phases. Renaming would have broken every existing pass subclass + pipeline. The framework-owned `attributes` field is skipped in `strategy.py`'s child walkers — it carries per-pass metadata, NOT IR children; `dataclasses.replace` preserves the field's value when not in kwargs so the resulting op still carries the same AttributeDict reference. Quality gates: - pytest tests/ — 1720 passed (1674 baseline + 46 new), 5 skipped - byte-equiv suite (test_runner_byte_equivalence + test_byte_equivalence_jit + test_cuda_complete_runner) — 532 passed, 2 skipped (matches baseline exactly) - mypy src/srdatalog/ — 148 errors in 40 files (matches baseline; zero new from this PR) - ruff check + ruff format — clean - D10 (LowerCtx pinned at 5 fields) — untouched See `docs/phase_decomposition_redesign.md` § 1 (ACID test + Amendments 1/2/3), § 3.2.1.2 (the 5 self-audit risks), § 3.2.1.3 (the pragma plugin contract), § 6.0.0 (PR-P0 scope), § 11 (MLIR / LLVM / Cascades / FRP precedent). Co-Authored-By: Claude Opus 4.7 (1M context) --- src/srdatalog/ir/core/__init__.py | 35 + src/srdatalog/ir/core/attribute.py | 201 ++++++ src/srdatalog/ir/core/dialect.py | 229 ++++++- src/srdatalog/ir/core/ops.py | 40 +- src/srdatalog/ir/core/passes.py | 117 +++- src/srdatalog/ir/core/pragma.py | 196 +++++- src/srdatalog/ir/core/pragma_plugin.py | 254 ++++++++ src/srdatalog/ir/core/services.py | 149 +++++ src/srdatalog/ir/core/strategy.py | 30 +- tests/test_pragma_plugin_framework.py | 856 +++++++++++++++++++++++++ 10 files changed, 2083 insertions(+), 24 deletions(-) create mode 100644 src/srdatalog/ir/core/attribute.py create mode 100644 src/srdatalog/ir/core/pragma_plugin.py create mode 100644 src/srdatalog/ir/core/services.py create mode 100644 tests/test_pragma_plugin_framework.py diff --git a/src/srdatalog/ir/core/__init__.py b/src/srdatalog/ir/core/__init__.py index c68dfa4..0733b78 100644 --- a/src/srdatalog/ir/core/__init__.py +++ b/src/srdatalog/ir/core/__init__.py @@ -42,6 +42,7 @@ from typing import NoReturn +from srdatalog.ir.core.attribute import Attribute, AttributeDict from srdatalog.ir.core.dialect import Compiler, Dialect from srdatalog.ir.core.lower_ctx import LowerCtx, NameGen, ViewLayout from srdatalog.ir.core.ops import Op, Type @@ -51,12 +52,14 @@ LoweringMissingError, LoweringPass, Pass, + PassCycleError, PassDriver, PassOrderingError, ProgramPass, Rewrite, RewritePass, program_pass, + topo_sort_passes, ) from srdatalog.ir.core.plugin import ( DIALECT_ENTRY_POINT_GROUP, @@ -68,6 +71,7 @@ PluginLoadError, ) from srdatalog.ir.core.pragma import ( + MaterializePragmaPass, Pragma, PragmaConfigError, PragmaCtx, @@ -78,7 +82,22 @@ has_pragma, pragma_handler, ) +from srdatalog.ir.core.pragma_plugin import ( + AttributeNameCollisionError, + MissingRequiredServiceError, + OpNameCollisionError, + PragmaPlugin, + PragmaPluginConflictError, + Render, + RenderDoubleRegistrationError, + pragma_plugin, +) from srdatalog.ir.core.scope import EmptyScope, Scope +from srdatalog.ir.core.services import ( + ServiceConflictError, + ServiceMissingError, + Services, +) from srdatalog.ir.core.strategy import ( Strategy, all_, @@ -116,6 +135,9 @@ def assert_never(value: object) -> NoReturn: 'ENTRY_POINT_GROUP', 'TARGET_ENTRY_POINT_GROUP', 'AmbiguousLowering', + 'Attribute', + 'AttributeDict', + 'AttributeNameCollisionError', 'Compiler', 'Dialect', 'EmptyScope', @@ -123,9 +145,13 @@ def assert_never(value: object) -> NoReturn: 'Lowering', 'LoweringMissingError', 'LoweringPass', + 'MaterializePragmaPass', + 'MissingRequiredServiceError', 'NameGen', 'Op', + 'OpNameCollisionError', 'Pass', + 'PassCycleError', 'PassDriver', 'PassOrderingError', 'PluginConflictError', @@ -136,10 +162,17 @@ def assert_never(value: object) -> NoReturn: 'PragmaConfigError', 'PragmaCtx', 'PragmaOrderingError', + 'PragmaPlugin', + 'PragmaPluginConflictError', 'ProgramPass', + 'Render', + 'RenderDoubleRegistrationError', 'Rewrite', 'RewritePass', 'Scope', + 'ServiceConflictError', + 'ServiceMissingError', + 'Services', 'Strategy', 'Type', 'UnconsumedPragmaError', @@ -157,11 +190,13 @@ def assert_never(value: object) -> NoReturn: 'id_', 'one', 'pragma_handler', + 'pragma_plugin', 'program_pass', 'repeat', 'seq', 'some', 'top_down', + 'topo_sort_passes', 'try_', 'verify_renderability', ] diff --git a/src/srdatalog/ir/core/attribute.py b/src/srdatalog/ir/core/attribute.py new file mode 100644 index 0000000..6c611e2 --- /dev/null +++ b/src/srdatalog/ir/core/attribute.py @@ -0,0 +1,201 @@ +'''Typed Op attributes — Risks 2 + 5 fix from spec § 3.2.1.2. + +Replaces "per-feature named fields on Op classes" with an MLIR-style +typed attribute dict. Adding a per-feature carrier (e.g. `ViewBinding`, +`OutputBinding`, `BgPlacementAttr`) is no longer a field-add on a +shared Op class — it's a new `Attribute` subclass packaged in a +`PragmaPlugin.new_attributes` tuple, looked up at the call site as +`op.attributes[ViewBinding]`. + +Per spec § 3.2.1.2: + + Risk 2: IR op classes grow per-feature with named binding fields + (`KernelDef.view_bindings`, `KernelDef.output_binding`, + ...). Each feature edits `KernelDef`. Bool-field anti- + pattern at the IR layer. + Fix: MLIR-style typed attribute dict on every Op: + `op.attributes[ViewBinding]`. Pragma plugins register + their own attribute types; Op classes don't grow. + + Risk 5: Op classes themselves carry fixed schemas (...) IS a + fixed-schema attribute carrier at the IR layer. + Fix: Ops are typed shells declaring only name + region + structure + verification. Per-feature data lives in + `op.attributes`. (Risk 2 fix generalized.) + +The `Attribute` base class is the pragma-side parallel of `Op` / +`Pragma` / `Type` (also pure data, frozen, slotted). The +`AttributeDict` storage is typed-key — each attribute kind is its own +TYPE, mirroring the `Services` shape (Risk 1 fix). + +This module is part of `core/`. Per D6 it knows no concrete dialects. + +See spec § 11.2 for the direct MLIR precedent. +''' + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, TypeVar + +# Bounded TypeVar so `AttributeDict[T]` only accepts Attribute +# subclasses. PEP 695 (`class AttributeDict: def __getitem__[T: ...]`) +# is Python 3.12+; the project targets 3.10, so we use the classic +# TypeVar form with a `bound` argument. The public API still types +# correctly under mypy --strict. + + +@dataclass(frozen=True, slots=True) +class Attribute: + '''Base class for typed attributes on Ops. + + Subclasses are typically `@final @dataclass(frozen=True, slots=True)`, + per the same discipline as `Pragma` (spec § 3.2.1.3) and `Op` (D1 / + D2 / D3 / D4 in `docs/code_discipline.md`). Each subclass declares + the per-feature data it carries; new features ship new Attribute + subclasses in a `PragmaPlugin.new_attributes` tuple. Op classes do + NOT grow per-feature fields. + + Usage (from a pragma plugin): + + @final + @dataclass(frozen=True, slots=True) + class ViewBinding(Attribute): + handle_idx: int + var_name: str + + # In the producing pass: + op.attributes[ViewBinding] = ViewBinding(handle_idx=0, var_name='v0') + + # In the consuming render: + vb = op.attributes[ViewBinding] + render(vb.var_name, vb.handle_idx) + + See spec § 3.2.1.2 (Risks 2 + 5). + ''' + + +_T = TypeVar('_T', bound=Attribute) + + +class AttributeDict: + '''Typed-key dict on every Op. `attributes[T] -> T instance`. + + Storage is `dict[type[Attribute], Attribute]` internally. The + method-level type variables surface mypy-correct retrieval at the + call site: `op.attributes[ViewBinding]` infers `ViewBinding`, not + `Any`. + + Mutation semantics: + - `__setitem__` overwrites silently. Attribute dicts are typically + populated by one producer pass; if two passes try to write the + same Attribute type, that is a coordination bug and the second + writer should be aware. We do NOT enforce conflict detection on + attribute writes — that's the producer's responsibility (each + producer pass should write its own Attribute types). See spec § + 3.2.1.3 for the plugin-contract conflict-detection that DOES + enforce ownership at registration time. + - `__delitem__` is intentionally not provided; per spec § 3.2.1.2 + Bucket 1, attributes ARE the cross-pass IR carrier, so deletion + is conceptually deleting IR data. The framework doesn't need it + yet; YAGNI. + + Discipline: + - Adding a new attribute kind = new `Attribute` subclass, registered + via a `PragmaPlugin.new_attributes` tuple. The `Op` class never + grows per-feature fields — all per-feature data lives here. + - Two pragma plugins MUST NOT register the same Attribute subclass. + `PragmaPlugin` registration detects the conflict atomically + (spec § 3.2.1.3); this dict trusts the registration layer. + + This object is mutable — it has to be, because passes populate it + incrementally. The owning `Op` is still frozen via `compare=False, + repr=False` on its `attributes` field (see `op.py`); the dict's + identity is excluded from op equality / hashing so two ops with the + same fields but different attribute states stay equal at the IR + level. Cross-pass data semantics live in the attributes; equality + semantics live in the op's frozen fields. + ''' + + __slots__ = ('_by_type',) + + def __init__(self) -> None: + self._by_type: dict[type[Attribute], Attribute] = {} + + def __getitem__(self, attr_type: type[_T]) -> _T: + '''Look up the Attribute instance keyed by `attr_type`. + + Raises `KeyError` (with the Attribute subclass name) if missing. + Callers that prefer a None fallback use `.get(...)` instead. + ''' + try: + return self._by_type[attr_type] # type: ignore[return-value] + except KeyError: + raise KeyError( + f'no Attribute of type {attr_type.__name__!r} on this op; ' + f'call sites that may legitimately have no instance should use ' + f'`op.attributes.get({attr_type.__name__})` instead.' + ) from None + + def __setitem__(self, attr_type: type[_T], value: _T) -> None: + '''Store `value` under the key `attr_type`. Silent overwrite — + see the class docstring for the conflict-discipline rationale. + + Validates that `value` is an instance of `attr_type` (cheap + guardrail against the `op.attributes[A] = b_instance` typo — + catches it at write time, not at the eventual read crash). + ''' + if not isinstance(value, attr_type): + raise TypeError( + f'AttributeDict[{attr_type.__name__}] requires an instance of ' + f'{attr_type.__name__}; got {type(value).__name__!r}.' + ) + self._by_type[attr_type] = value + + def __contains__(self, attr_type: object) -> bool: + '''Membership test: `ViewBinding in op.attributes`. + + The argument is typed `object` so callers can write the natural + `if X in op.attributes:` without mypy complaining about generic- + parameter inference. Production lookups go through `__getitem__` + / `get(...)`; this is the cheap "do you have one?" path. + ''' + return attr_type in self._by_type + + def get(self, attr_type: type[_T]) -> _T | None: + '''Look up the Attribute instance keyed by `attr_type`, or None. + + Symmetric with `dict.get(key, None)`. Never raises. + ''' + val = self._by_type.get(attr_type) + if val is None: + return None + # Defensive: a misregistered attribute under the wrong key would + # otherwise mistype this return; surface it loudly. + if not isinstance(val, attr_type): # pragma: no cover (registration-time guard) + raise TypeError( + f'AttributeDict.get({attr_type.__name__}) found a {type(val).__name__!r} ' + f'under the {attr_type.__name__!r} key. Registration discipline broken; ' + f'see spec § 3.2.1.3 for the plugin-contract conflict-detection rules.' + ) + return val + + def __len__(self) -> int: + return len(self._by_type) + + def __iter__(self) -> Any: + '''Iterate over registered Attribute TYPES. Mirrors `dict.__iter__` + which yields keys.''' + return iter(self._by_type) + + def __repr__(self) -> str: + if not self._by_type: + return 'AttributeDict()' + inner = ', '.join(f'{k.__name__}={v!r}' for k, v in self._by_type.items()) + return f'AttributeDict({inner})' + + +__all__ = [ + 'Attribute', + 'AttributeDict', +] diff --git a/src/srdatalog/ir/core/dialect.py b/src/srdatalog/ir/core/dialect.py index 11f12b3..e372a57 100644 --- a/src/srdatalog/ir/core/dialect.py +++ b/src/srdatalog/ir/core/dialect.py @@ -31,10 +31,12 @@ from __future__ import annotations +import threading import warnings -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager from dataclasses import dataclass, field -from typing import Any +from typing import TYPE_CHECKING, Any from srdatalog.ir.core.plugin import ( DIALECT_ENTRY_POINT_GROUP, @@ -47,6 +49,13 @@ _info_for, _topo_sort_plugins, ) +from srdatalog.ir.core.services import Services + +if TYPE_CHECKING: + from srdatalog.ir.core.attribute import Attribute + from srdatalog.ir.core.ops import Op + from srdatalog.ir.core.passes import Lowering, Pass + from srdatalog.ir.core.pragma_plugin import PragmaPlugin, Render # One-shot DeprecationWarning emission marker. PR-1a's spec # (`docs/phase_decomposition_redesign.md` §3.3.4) requires the legacy @@ -89,6 +98,56 @@ class Dialect: _DIRECT_OWNER = '' +# PR-P0 (spec § 3.2.1.2 Risk 4 fix): per-Compiler "current compiler" +# thread-local. Decorator sugar (@register_render / @lowering / +# @pragma_handler) consults this; when set, the decorator stages its +# registration into the per-Compiler registries; otherwise it falls +# back to the legacy module-global registries (back-compat shim, +# deleted in PR-P5). +# +# This dissolves the Risk-4 anti-pattern ("decorators wrap mutable +# module-global state") without requiring every existing decorator +# call site to grow a `compiler=` kwarg. The thread-local is set by +# `Compiler.register_pragma_plugin(...)` while importing / running +# the plugin's contribution code; outside that scope the decorators +# preserve their existing module-global behavior. +_current_compiler: threading.local = threading.local() + + +def _get_current_compiler() -> Compiler | None: + '''Return the thread-local "current compiler", or None if not in a + `_compiler_registration_scope`. Decorator sugar reads this to + decide whether to stage into the per-Compiler registries or fall + back to module-globals. + ''' + return getattr(_current_compiler, 'compiler', None) + + +@contextmanager +def _compiler_registration_scope(compiler: Compiler) -> Iterator[None]: + '''Set the thread-local "current compiler" for the duration of the + with-block (PR-P0, spec § 3.2.1.2 Risk 4 fix). + + Used by `Compiler.register_pragma_plugin` to make decorators stage + registrations into THIS compiler's per-Compiler registries instead + of the legacy module-globals. Nested scopes are tolerated and + restore the previous value on exit (sequential plugin loads from + inside another plugin's `register` are legal). + ''' + import contextlib + + prev = getattr(_current_compiler, 'compiler', None) + _current_compiler.compiler = compiler + try: + yield + finally: + if prev is None: + with contextlib.suppress(AttributeError): + del _current_compiler.compiler + else: + _current_compiler.compiler = prev + + class Compiler: '''Holds the registered dialects and tracks which plugin owns each. @@ -125,6 +184,30 @@ def __init__(self) -> None: # See docs/phase_decomposition_redesign.md §3.3.1. self.active_target: str = 'cuda' + # ---------------------------------------------------------------- + # PR-P0 per-Compiler registries (spec § 3.2.1.2 Risk 4 fix + + # § 3.2.1.3). The legacy module-globals (`_PRAGMA_REGISTRY`, + # `_PLUGIN_REGISTRY`, the cuda render registries) stay around as + # back-compat shims through PR-P1..P4; PR-P5 deletes them. From + # PR-P0 onward these per-Compiler dicts are the SOURCE OF TRUTH + # for everything the decorator sugar stages while a + # `_compiler_registration_scope(self)` is active. + # ---------------------------------------------------------------- + self.ops: dict[str, type[Op]] = {} + self.attributes: dict[str, type[Attribute]] = {} + self.passes: list[Pass] = [] # topo-sorted; see register_pragma_plugin + # Lowerings keyed on (source_op_type, target_op_type). Today's + # `Lowering.matches` is `type` (un-bound, pre-existing field + # shape from passes.py); the typed-by-op narrowing happens in + # PR-P1..P4 as each plugin migrates. The dict is typed `type`, + # not `type[Op]`, to stay compatible with the un-narrowed Lowering + # shape while still encoding the (source, target) pair shape from + # spec § 6.0.0. + self.lowerings: dict[tuple[type, type], Lowering] = {} + self.renders: dict[tuple[type, str], Render] = {} + self.services: Services = Services() + self.pragma_plugins: list[PragmaPlugin] = [] + # -- dialect registration -------------------------------------------------- def register_dialect(self, d: Dialect) -> None: @@ -181,10 +264,23 @@ def dialects(self) -> list[Dialect]: '''All registered dialects, in registration order.''' return list(self._dialects.values()) - def run(self, prog: Any, *, pipeline: list[Any], target: str = 'cuda') -> Any: + def run( + self, + prog: Any, + *, + pipeline: list[Any] | None = None, + target: str = 'cuda', + ) -> Any: '''Drive a list of `Pass` instances over `prog`. Returns the (possibly transformed) prog. + PR-P0: `pipeline` is now optional. If omitted, the run uses + `self.passes` (the topo-sorted per-Compiler pass list maintained + by `register_pragma_plugin(...)`). Existing call sites that pass + `pipeline=[...]` continue to work unchanged — the existing + explicit-pipeline construction surface stays around through the + per-pragma migration phases. + Pre-flight: validate pipeline ordering. For each `Pass`, check that all `consumes` dialects are either registered with this Compiler OR produced by an earlier Pass in the list. Raises @@ -205,8 +301,10 @@ def run(self, prog: Any, *, pipeline: list[Any], target: str = 'cuda') -> Any: # (passes.py imports Compiler from this module). from srdatalog.ir.core.passes import PassOrderingError + pipeline_list: list[Any] = pipeline if pipeline is not None else list(self.passes) + available: set[str] = {d.name for d in self.dialects} - for i, p in enumerate(pipeline): + for i, p in enumerate(pipeline_list): for needed in p.consumes: if needed not in available: raise PassOrderingError(p.name, needed, i) @@ -215,7 +313,7 @@ def run(self, prog: Any, *, pipeline: list[Any], target: str = 'cuda') -> Any: prev_target = self.active_target self.active_target = target try: - for p in pipeline: + for p in pipeline_list: prog = p.apply(prog, self) return prog finally: @@ -275,6 +373,127 @@ def register_plugin(self, plugin: Any) -> None: self._plugins_loaded[plugin_name] = info + # -- PR-P0: atomic pragma-plugin registration ----------------------------- + + def register_pragma_plugin(self, plugin: PragmaPlugin) -> None: + '''Atomically unpack a `PragmaPlugin` into the per-Compiler + registries (PR-P0, spec § 3.2.1.3 / § 6.0.0). + + All-or-nothing transaction. Conflicts detected up front; if any + is found, the per-Compiler registries are left UNCHANGED — no + partial registration. Conflicts checked (spec § 6.0.0): + + - Op-name collision: `plugin.new_ops` contains an op whose + `__name__` is already in `self.ops` under a different class. + - Attribute-name collision: same shape on `plugin.new_attributes` + vs `self.attributes`. + - Render double-registration: `(op_type, target)` already in + `self.renders`. + - Pass-cycle: after appending `plugin.passes`, the LLVM-style + topo-sort raises `PassCycleError`. + - Missing required service: a type in + `plugin.requires_services` is not registered with + `self.services`. + + On success: + + - `self.ops` / `self.attributes` / `self.renders` populated + with the plugin's contributions. + - `self.lowerings` populated keyed on `(matches, target_op)` + pairs. Today's `Lowering` carries only `matches` (source op + type); the target op is read from a single produced op when + unambiguous. Multi-target lowerings are a Phase-B2 concern + outside PR-P0 scope; for now, single-target lowerings store + under `(matches, matches)` if no `produces_ops` is declared. + - `self.passes` extended with the plugin's passes, then + topo-sorted via `topo_sort_passes(...)`. + - `self.pragma_plugins` appended for introspection. + + See spec § 3.2.1.3 for the rationale ("one atomic registration + that bundles ALL of [pragma_cls, ops, attributes, passes, + lowerings, renders, services, dependencies]"). + ''' + # Local imports avoid the core-init cycle. pragma_plugin imports + # passes which imports dialect (this module). + from srdatalog.ir.core.passes import PassCycleError, topo_sort_passes + from srdatalog.ir.core.pragma_plugin import ( + AttributeNameCollisionError, + MissingRequiredServiceError, + OpNameCollisionError, + RenderDoubleRegistrationError, + ) + + # -------- Phase 1: conflict detection (no mutation) -------- + + # Op-name collisions. + for op_cls in plugin.new_ops: + existing = self.ops.get(op_cls.__name__) + if existing is not None and existing is not op_cls: + raise OpNameCollisionError( + f'PragmaPlugin {plugin.pragma_cls.__name__!r}: op name ' + f'{op_cls.__name__!r} already registered to {existing!r}.' + ) + + # Attribute-name collisions. + for attr_cls in plugin.new_attributes: + existing_attr = self.attributes.get(attr_cls.__name__) + if existing_attr is not None and existing_attr is not attr_cls: + raise AttributeNameCollisionError( + f'PragmaPlugin {plugin.pragma_cls.__name__!r}: attribute name ' + f'{attr_cls.__name__!r} already registered to {existing_attr!r}.' + ) + + # Render double-registration. + for target, render_tuple in plugin.renders.items(): + for r in render_tuple: + key = (r.op_type, target) + if key in self.renders: + raise RenderDoubleRegistrationError( + f'PragmaPlugin {plugin.pragma_cls.__name__!r}: render for ' + f'(op_type={r.op_type.__name__!r}, target={target!r}) is ' + f'already registered.' + ) + + # Missing required services. + for svc_type in plugin.requires_services: + if svc_type not in self.services: + raise MissingRequiredServiceError( + f'PragmaPlugin {plugin.pragma_cls.__name__!r}: requires service ' + f'of type {svc_type.__name__!r}, but none is registered with ' + f'this Compiler. Register the service before this plugin (e.g. ' + f'`compiler.services.register({svc_type.__name__}, )`).' + ) + + # Pass-cycle detection. Build the *proposed* combined pass list + # and topo-sort it in a staging variable; commit only on success. + proposed_passes = [*self.passes, *plugin.passes] + try: + sorted_passes = topo_sort_passes(proposed_passes) + except PassCycleError as e: + raise PassCycleError(e.cycle) from None + + # -------- Phase 2: commit (no early returns past this point) -------- + + for op_cls in plugin.new_ops: + self.ops[op_cls.__name__] = op_cls + + for attr_cls in plugin.new_attributes: + self.attributes[attr_cls.__name__] = attr_cls + + for low in plugin.lowerings: + # Single-target lowering: key on (matches, matches) when no + # explicit target-op is known. Multi-target dispatching is + # Phase B2 work. The key shape mirrors spec § 6.0.0 's + # `lowerings: dict[tuple[type[Op], type[Op]], Lowering]`. + self.lowerings[(low.matches, low.matches)] = low + + for target, render_tuple in plugin.renders.items(): + for r in render_tuple: + self.renders[(r.op_type, target)] = r + + self.passes = sorted_passes + self.pragma_plugins.append(plugin) + @staticmethod def _resolve_plugin(plugin: Any) -> tuple[str, Callable[[Any], None]]: '''Coerce a plugin argument to `(name, register_fn)`. diff --git a/src/srdatalog/ir/core/ops.py b/src/srdatalog/ir/core/ops.py index 9166162..db67172 100644 --- a/src/srdatalog/ir/core/ops.py +++ b/src/srdatalog/ir/core/ops.py @@ -2,7 +2,8 @@ `Op` and `Type` are pure-data marker bases. Subclasses are dataclasses with `frozen=True, slots=True`. The bases carry no methods and no -state. +state — *with the single exception of `Op.attributes`*, the typed +attribute dict introduced in PR-P0 (spec § 3.2.1.2 Risks 2 + 5). Discipline rules (see docs/design_principles.md): @@ -12,11 +13,31 @@ - Dispatch via `match`, not method override. These are enforced by tests/test_ir_core_discipline.py. + +About `Op.attributes` (PR-P0, spec § 3.2.1.2 / § 3.2.1.3): + + - Default factory: a fresh empty `AttributeDict` per instance. + - `kw_only=True`: keeps positional construction stable for every + existing subclass that already declares its own positional + fields. Without this, adding any defaulted parent field would + break every subclass with non-defaulted fields (Python's + "non-default arg follows default arg" rule on the generated + `__init__`). + - `compare=False, hash=False`: attribute state is per-pass metadata, + NOT part of structural op identity. Two ops with the same frozen + fields are equal regardless of attribute state — equality lives + in the IR shape, cross-pass data lives in attributes (spec § + 3.2.1.3, MLIR precedent § 11.2). + - `repr=False`: keeps the repr clean and byte-equivalent with the + pre-PR-P0 shape; byte-equivalence goldens that print ops are + unaffected. ''' from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field + +from srdatalog.ir.core.attribute import AttributeDict @dataclass(frozen=True, slots=True) @@ -37,4 +58,19 @@ class Op: Each dialect defines its own subclasses, typically as dataclasses. Op instances form an IR program; the pass driver walks them via pattern matching, never via virtual dispatch on methods. + + `attributes` is the MLIR-style typed attribute dict (PR-P0, spec § + 3.2.1.2 Risks 2 + 5). Per-feature data — view bindings, output + bindings, index bindings, block-group placement, etc. — lives there, + registered by the owning pragma plugin via + `PragmaPlugin.new_attributes`. Op subclasses do NOT grow per-feature + fields. See `attribute.py` for the storage type. ''' + + attributes: AttributeDict = field( + default_factory=AttributeDict, + kw_only=True, + compare=False, + hash=False, + repr=False, + ) diff --git a/src/srdatalog/ir/core/passes.py b/src/srdatalog/ir/core/passes.py index d05615a..43f09eb 100644 --- a/src/srdatalog/ir/core/passes.py +++ b/src/srdatalog/ir/core/passes.py @@ -373,15 +373,37 @@ class Pass(ABC): Per `docs/compiler_redesign.md` §4: every step in every `Compiler.run` pipeline is one of these three kinds; nothing else. - `consumes` / `produces` declare dialect dependencies for the - ordering validation done by `Compiler.run` at construction time - (per the R4 research report and `phase_b_lowering_dispatcher.md` - §5). They are advisory data, not enforced inside `apply`. + `consumes` / `produces` (str-typed, dialect names) declare dialect + dependencies for the dialect-level ordering validation done by + `Compiler.run` at construction time (per the R4 research report and + `phase_b_lowering_dispatcher.md` §5). They are advisory data, not + enforced inside `apply`. + + `consumes_ops` / `produces_ops` / `preserves` (PR-P0, spec § + 3.2.1.3) declare op-type-level dependencies for the LLVM-style + topological scheduler. A pass that `produces_ops=(BgRoot, ...)` + must run before any pass that `consumes_ops=(BgRoot, ...)`. The + scheduler (`topo_sort_passes`) is consulted by + `Compiler.register_pragma_plugin` when a plugin contributes new + passes; the existing dialect-name `consumes` / `produces` remain + the surface for legacy pipelines built by hand. + + Naming note: the spec example in § 3.2.1.3 / § 6.0.0 lists Pass + fields as `produces` / `consumes` (typed as op classes). The + pre-PR-P0 framework already uses those names for str-typed dialect + IDs (`tuple[str, ...]`); renaming would break every existing pass + subclass and every pipeline that constructs them. The op-typed + fields are introduced under suffixed names (`produces_ops`, + `consumes_ops`, `preserves`) so the two ordering surfaces coexist + cleanly during the per-pragma migration phases (PR-P1..P5). ''' name: str consumes: tuple[str, ...] = () produces: tuple[str, ...] = () + consumes_ops: tuple[type, ...] = () + produces_ops: tuple[type, ...] = () + preserves: tuple[type, ...] = () @abstractmethod def apply(self, prog: Any, compiler: Any) -> Any: @@ -555,6 +577,91 @@ def apply(self, prog: Any, compiler: Any) -> Any: return self.fn(prog, compiler) +class PassCycleError(Exception): + '''Cycle in the op-type-level produces/consumes graph across the + passes registered with a `Compiler` (PR-P0, spec § 3.2.1.3). + + Raised by `topo_sort_passes` when the LLVM-style scheduler cannot + produce a linear order. Carries the cycle as a list of pass names + for diagnostics. + + Distinct from `PassOrderingError` (pipeline-position-based dialect + dependency error). `PassCycleError` is the op-type-graph error + raised during plugin registration; `PassOrderingError` is the + pipeline-construction error raised at `Compiler.run` time. + ''' + + def __init__(self, cycle: list[str]) -> None: + self.cycle = cycle + super().__init__( + f'pass produces/consumes graph has a cycle involving: {cycle}. ' + f'A pass that consumes an op type must run after every pass that ' + f'produces it; declare an alternative ordering via `preserves` or ' + f'split the cyclic responsibility across separate passes. See ' + f'spec § 3.2.1.3 (LLVM-style topo-sort).' + ) + + +def topo_sort_passes(passes: list[Pass]) -> list[Pass]: + '''Topologically sort `passes` by op-type-level produces/consumes. + + LLVM legacy-pass-manager shape (spec § 11.3): a pass that + `produces_ops=(X, ...)` must run before any pass that + `consumes_ops=(X, ...)`. Passes that share neither relation may run + in any order; the scheduler picks a deterministic order (stable on + the input list, lex-tiebreak on pass name when sets allow it). + + Returns a new list of `Pass` instances in topo-sorted order. Raises + `PassCycleError` if no linear order exists. + + Self-edges (a pass that both produces and consumes the same op + type) are tolerated — the produces side is interpreted as "after + this pass, instances of X exist", and the consumes side as "this + pass may read X". A pass legitimately may produce and consume the + same type (the materialization passes do exactly that). + ''' + # Map op type -> indices of passes that produce that type. + producers: dict[type, list[int]] = {} + for i, p in enumerate(passes): + for op_type in p.produces_ops: + producers.setdefault(op_type, []).append(i) + + # Build edges: producer index -> consumer index. + out_edges: dict[int, set[int]] = {i: set() for i in range(len(passes))} + in_degree: dict[int, int] = dict.fromkeys(range(len(passes)), 0) + for j, p in enumerate(passes): + for op_type in p.consumes_ops: + for i in producers.get(op_type, []): + if i == j: + # Self-edge tolerated; see docstring. + continue + if j not in out_edges[i]: + out_edges[i].add(j) + in_degree[j] += 1 + + # Kahn: pop lex-smallest (by pass name, then original index) zero- + # in-degree node, repeat. Stable on inputs where the graph permits. + def _sort_key(i: int) -> tuple[str, int]: + return (passes[i].name, i) + + ready = sorted([i for i, d in in_degree.items() if d == 0], key=_sort_key) + order: list[int] = [] + while ready: + i = ready.pop(0) + order.append(i) + new_ready: list[int] = [] + for j in sorted(out_edges[i], key=_sort_key): + in_degree[j] -= 1 + if in_degree[j] == 0: + new_ready.append(j) + ready = sorted(set(ready) | set(new_ready), key=_sort_key) + + if len(order) != len(passes): + remaining = sorted(set(range(len(passes))) - set(order)) + raise PassCycleError([passes[i].name for i in remaining]) + return [passes[i] for i in order] + + def program_pass( *, name: str, @@ -594,6 +701,7 @@ def _wrap(fn: Callable[[Any, Any], Any]) -> ProgramPass: 'LoweringMissingError', 'LoweringPass', 'Pass', + 'PassCycleError', 'PassDependencyError', 'PassDriver', 'PassOrderingError', @@ -603,5 +711,6 @@ def _wrap(fn: Callable[[Any, Any], Any]) -> ProgramPass: 'lowering', 'program_pass', 'rewrite', + 'topo_sort_passes', 'verifier', ] diff --git a/src/srdatalog/ir/core/pragma.py b/src/srdatalog/ir/core/pragma.py index 4da5358..64bb44d 100644 --- a/src/srdatalog/ir/core/pragma.py +++ b/src/srdatalog/ir/core/pragma.py @@ -265,15 +265,36 @@ def handler(op: Op, pragma: pragma_cls, ctx: PragmaCtx) -> Op | None: def _wrap( fn: Callable[[Any, Pragma, PragmaCtx], Any], ) -> Callable[[Any, Pragma, PragmaCtx], Any]: - _PRAGMA_REGISTRY.append( - PragmaRegistration( - pragma_cls=pragma_cls, - on=on, - fn=fn, - before=tuple(before), - after=tuple(after), - ) + reg = PragmaRegistration( + pragma_cls=pragma_cls, + on=on, + fn=fn, + before=tuple(before), + after=tuple(after), ) + + # PR-P0 back-compat shim (spec § 6.0.0 row 8): if a + # `_compiler_registration_scope(compiler)` is active, stage the + # registration into THAT compiler's per-Compiler pragma_handlers + # list. Otherwise fall back to the legacy module-global registry + # (existing behavior, deleted in PR-P5). + # + # The per-Compiler list lives on `compiler._pragma_handlers` + # (private attr; the framework owns it). It is populated either + # directly here (when this decorator runs inside a registration + # scope) or by `register_pragma_plugin` unpacking the plugin's + # `pragma_cls` into a synthetic legacy registration. PR-P5 + # collapses both code paths into the typed PragmaPlugin contract. + from srdatalog.ir.core.dialect import _get_current_compiler + + compiler = _get_current_compiler() + if compiler is not None: + handlers: list[PragmaRegistration] = getattr(compiler, '_pragma_handlers', []) + if not handlers: + compiler._pragma_handlers = handlers # type: ignore[attr-defined] + handlers.append(reg) + else: + _PRAGMA_REGISTRY.append(reg) return fn return _wrap @@ -327,7 +348,165 @@ def has_pragma(op: Any, pragma_cls: type[Pragma]) -> bool: return any(isinstance(p, pragma_cls) for p in pragmas) +# ----------------------------------------------------------------------------- +# PR-P0 back-compat Pass: MaterializePragmaPass +# +# When a `@pragma_handler` is staged into a per-Compiler scope via +# `_compiler_registration_scope`, the framework needs a Pass instance +# that actually runs those handlers as part of `compiler.run(...)`. +# Today's `mir.pragma_pass.MirPragmaPass` is the production handler- +# runner, but it pulls from the module-global registry. The shim below +# is a lightweight Pass that iterates `get_compiler_pragma_handlers` +# (which returns the per-Compiler list when present) and applies each +# handler to matching ops in `prog`. +# +# PR-P5 deletes this shim alongside the module-global registry; by +# then every pragma is a full `PragmaPlugin` with its own typed Pass. +# ----------------------------------------------------------------------------- + + +def _materialize_pragma_pass_factory() -> Any: + '''Lazily build the `MaterializePragmaPass` class so its + dependency on `Pass` (from `passes.py`, which imports Compiler from + `dialect.py`) doesn't introduce a module-load cycle. + + Returns the class; instances are constructed by callers. + ''' + from dataclasses import dataclass as _dataclass + from dataclasses import field as _field + + from srdatalog.ir.core.passes import Pass + + @_dataclass(frozen=True) + class MaterializePragmaPass(Pass): + '''PR-P0 back-compat shim Pass: materialize one Pragma subclass. + + Runs the registered `@pragma_handler(pragma_cls, on=)` + for `pragma_cls` against every matching op in `prog`. Hand-rolled + tree walk (no strategy combinator) so the shim has no dependency + on the IIR-walk infrastructure. + + Each handler signature is `(op, pragma, ctx) -> Op | None`. A + None return is treated as "this pragma did not apply to this op" + and leaves the op unchanged; a returned Op replaces the matched + op in the program tree. + ''' + + name: str = 'materialize_pragma_shim' + pragma_cls: type[Pragma] | None = _field(default=None) + + def apply(self, prog: Any, compiler: Any) -> Any: + handlers = get_compiler_pragma_handlers(compiler) + if self.pragma_cls is not None: + handlers = [h for h in handlers if h.pragma_cls is self.pragma_cls] + if not handlers: + return prog + + ctx = PragmaCtx(compiler=compiler) + + def _apply_handlers_to(op: Any) -> Any: + for h in handlers: + if not isinstance(op, h.on): + continue + pragmas = getattr(op, 'pragmas', None) + if pragmas is None: + continue + for p in pragmas: + if isinstance(p, h.pragma_cls): + new = h.fn(op, p, ctx) + if new is not None: + op = new + return op + + # Walk the tree top-down: visit the root, then recurse into + # dataclass fields that are Ops or contain Ops. Mirrors the + # MIR pragma-pass's reach without depending on its dialect + # specifics. + def _walk(node: Any) -> Any: + import dataclasses as _dc + + node = _apply_handlers_to(node) if _is_op_like(node) else node + if not _dc.is_dataclass(node) or isinstance(node, type): + return node + new_fields: dict[str, Any] = {} + any_changed = False + for f in _dc.fields(node): + val = getattr(node, f.name) + new_val = _walk_value(val) + new_fields[f.name] = new_val + if new_val is not val: + any_changed = True + if any_changed: + return _dc.replace(node, **new_fields) + return node + + def _walk_value(val: Any) -> Any: + if _is_op_like(val): + return _walk(val) + if isinstance(val, list): + new_list = [_walk_value(x) for x in val] + return new_list if any(a is not b for a, b in zip(new_list, val)) else val + if isinstance(val, tuple): + new_tuple = tuple(_walk_value(x) for x in val) + return new_tuple if any(a is not b for a, b in zip(new_tuple, val)) else val + return val + + def _is_op_like(node: Any) -> bool: + # An Op subclass instance; avoid importing Op at module init + # by walking MRO names. + cls = type(node) + for base in cls.__mro__: + if base.__name__ == 'Op' and base.__module__ == 'srdatalog.ir.core.ops': + return True + return False + + return _walk(prog) + + return MaterializePragmaPass + + +def get_compiler_pragma_handlers(compiler: Any) -> list[PragmaRegistration]: + '''PR-P0 back-compat shim (spec § 6.0.0 row 8): snapshot the + per-Compiler `_pragma_handlers` list staged by + `@pragma_handler` decorations that ran inside a + `_compiler_registration_scope(compiler)` block. + + Returns the legacy module-global registrations IF the compiler has + no per-instance staging (every existing test path), so callers + written before the per-Compiler shape continue to work unchanged. + Otherwise returns the per-Compiler staged registrations + concatenated with the module-globals (the compiler instance is + authoritative; module-globals are fallback for legacy code paths). + + Used by `MaterializePragmaPass` (a back-compat shim Pass that + iterates these registrations and applies them). After PR-P5, + pragmas register exclusively via `PragmaPlugin` and this helper + is deleted alongside the module-global registry. + ''' + per_compiler: list[PragmaRegistration] = list(getattr(compiler, '_pragma_handlers', [])) + if per_compiler: + return per_compiler + list(_PRAGMA_REGISTRY) + return list(_PRAGMA_REGISTRY) + + +def MaterializePragmaPass(*args: Any, **kwargs: Any) -> Any: + '''Public constructor for the back-compat shim Pass (PR-P0). + + Calls `_materialize_pragma_pass_factory()` on first use to avoid + the module-load cycle with `passes.py`; subsequent calls reuse the + cached class. + ''' + global _MATERIALIZE_PRAGMA_PASS_CLS + if _MATERIALIZE_PRAGMA_PASS_CLS is None: + _MATERIALIZE_PRAGMA_PASS_CLS = _materialize_pragma_pass_factory() + return _MATERIALIZE_PRAGMA_PASS_CLS(*args, **kwargs) + + +_MATERIALIZE_PRAGMA_PASS_CLS: type | None = None + + __all__ = [ + 'MaterializePragmaPass', 'Pragma', 'PragmaConfigError', 'PragmaCtx', @@ -335,6 +514,7 @@ def has_pragma(op: Any, pragma_cls: type[Pragma]) -> bool: 'PragmaRegistration', 'UnconsumedPragmaError', 'UnregisteredPragmaError', + 'get_compiler_pragma_handlers', 'get_pragma', 'get_pragma_registrations', 'has_pragma', diff --git a/src/srdatalog/ir/core/pragma_plugin.py b/src/srdatalog/ir/core/pragma_plugin.py new file mode 100644 index 0000000..08a4e08 --- /dev/null +++ b/src/srdatalog/ir/core/pragma_plugin.py @@ -0,0 +1,254 @@ +'''PragmaPlugin schema — the atomic pragma-as-plugin contract. + +PR-P0 / spec § 3.2.1.3: a pragma is packaged as one `PragmaPlugin` +that bundles every artifact the pragma needs — its `Pragma` subclass, +any new IR ops, any new typed attributes, the passes that materialize ++ optimize it, the lowerings + per-target renders. The whole bundle +registers atomically with one `compiler.register_pragma_plugin(...)` +call; partial state never leaks. + +This is the ACID-test fix (spec § 1): "adding a new pragma = new +plugin module, ZERO edits to existing source files." Today's +`@pragma_handler` cannot ship new ops, cannot bundle lowerings, and +cannot declare dependencies — see the five "what +`@pragma_handler` CANNOT do" bullets in spec § 3.2.1.3. + +Schema sketch (spec § 3.2.1.3): + + @final + @dataclass(frozen=True, slots=True) + class BlockGroupPlugin(PragmaPlugin): + pragma_cls: type[Pragma] = BlockGroupPragma + new_ops: tuple[type[Op], ...] = (BgRoot, ...) + new_attributes: tuple[type[Attribute], ...] = (BgPlacementAttr,) + passes: tuple[Pass, ...] = (MaterializeBlockGroup(),) + lowerings: tuple[Lowering, ...] = (LowerBgRoot(), ...) + renders: dict[str, tuple[Render, ...]] = field(default_factory=lambda: { + 'cuda': (RenderBgRootCuda(), ...), + }) + requires_services: tuple[type, ...] = (NameGen, ViewLayout) + produces_ops: tuple[type[Op], ...] = (BgRoot, ...) + consumes_ops: tuple[type[Op], ...] = (mir.ExecutePipeline,) + preserves: tuple[type, ...] = (semi_naive_safety, pragma_compat) + +The `register_pragma_plugin(plugin)` API on `Compiler` (spec § 6.0.0) +atomically unpacks all of these into the per-Compiler registries; +conflict detection on op-name, attribute-name, render-double- +registration, pass-cycle, and missing required service is part of +that path. See `dialect.py` for the unpack logic. + +Module discipline: per D6, this file lives in `core/` and imports no +concrete dialect; the plugin schema is dialect-agnostic. The +schema's fields are typed-shells (`type[Op]`, `type[Attribute]`, +`type[Pragma]`) — every concrete subclass lives in its owning plugin +module. +''' + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any, final + +from srdatalog.ir.core.attribute import Attribute +from srdatalog.ir.core.ops import Op +from srdatalog.ir.core.passes import Lowering, Pass +from srdatalog.ir.core.pragma import Pragma + + +@dataclass(frozen=True) +class Render: + '''A per-target render registration packaged inside a `PragmaPlugin`. + + A render produces target-specific output text (CUDA C++ today; TBB + / SYCL / ... post-Phase B2) for one Op type. Bundled into a plugin + so the plugin owns ALL of its renders — adding a new render does + not edit any framework file. + + Fields: + op_type — Op subclass this render handles (one render per + (op_type, target) pair; conflict detection in + `register_pragma_plugin` enforces uniqueness). + fn — render callable. Signature is target-specific (CUDA + today uses `(op, EmitCtx) -> str`; per-target details + live in the target's own contract). Framework-side + this is `Callable[..., Any]`. + mode — render dispatch mode. CUDA distinguishes 'stmt' vs + 'expr'; other targets may ignore. Defaults to 'stmt' + for back-compat with the existing CUDA registry. + + Per spec § 3.2.1.3, the `renders` field on `PragmaPlugin` is a + `dict[str, tuple[Render, ...]]` keyed on target name; one + `Render` is one (op_type, mode) registration for that target. + ''' + + op_type: type + fn: Callable[..., Any] + mode: str = 'stmt' + + +@dataclass(frozen=True) +class PragmaPlugin: + '''A pragma packaged as a compiler plugin. + + Bundles every artifact a pragma needs into one atomic registration. + See spec § 3.2.1.3 for the full contract and rationale. + + Construct directly: + + plugin = PragmaPlugin( + pragma_cls=MyPragma, + new_ops=(MyRoot, MyScope), + passes=(MaterializeMy(),), + lowerings=(LowerMyRoot(),), + renders={'cuda': (Render(MyRoot, render_my_root_cuda),)}, + produces_ops=(MyRoot, MyScope), + consumes_ops=(MirExecutePipeline,), + ) + compiler.register_pragma_plugin(plugin) + + Or build subclasses with `@final @dataclass(frozen=True, slots=True)` + (matching the same discipline as `Pragma` / `Op` — pure data). + Subclasses should use `field(default=...)` or `field(default_factory=...)` + for any pre-populated field, exactly mirroring the spec § 3.2.1.3 + example. + + Atomicity: `Compiler.register_pragma_plugin(plugin)` is all-or-nothing. + Conflicts detected up front (op-name collision, attribute-name + collision, render double-registration, pass-cycle, missing required + service) abort the registration with NO partial state in the per- + Compiler registries. See `dialect.py:register_pragma_plugin`. + + Discipline: + - `slots=True` is INTENTIONALLY omitted on `PragmaPlugin` itself + (not on its subclasses, which should still use it). The base + uses field defaults heavily; slots+inheritance with mixed + defaults pushes Python's dataclass machinery into edge cases + that don't pay back in this code path (the schema is constructed + at plugin-load time, not in a hot loop). + - frozen=True so a registered plugin object is immutable after + handoff; passes that capture a plugin reference cannot mutate + its bundle. + ''' + + pragma_cls: type[Pragma] + new_ops: tuple[type[Op], ...] = () + new_attributes: tuple[type[Attribute], ...] = () + passes: tuple[Pass, ...] = () + lowerings: tuple[Lowering, ...] = () + renders: dict[str, tuple[Render, ...]] = field(default_factory=dict) + requires_services: tuple[type, ...] = () + produces_ops: tuple[type[Op], ...] = () + consumes_ops: tuple[type[Op], ...] = () + preserves: tuple[type, ...] = () + + +def pragma_plugin( + pragma_cls: type[Pragma], + *, + new_ops: tuple[type[Op], ...] = (), + new_attributes: tuple[type[Attribute], ...] = (), + passes: tuple[Pass, ...] = (), + lowerings: tuple[Lowering, ...] = (), + renders: dict[str, tuple[Render, ...]] | None = None, + requires_services: tuple[type, ...] = (), + produces_ops: tuple[type[Op], ...] = (), + consumes_ops: tuple[type[Op], ...] = (), + preserves: tuple[type, ...] = (), +) -> Callable[[type], PragmaPlugin]: + '''Class-decorator factory: build a `PragmaPlugin` from a class body. + + Sugar over the bare `PragmaPlugin(pragma_cls=..., ...)` constructor. + Symmetric with `@pragma_handler` / `@lowering` / `@register_render` + — the decorator IS the registration shape (Triton-style; see + `feedback_decorator_registries.md`). + + Usage: + + @pragma_plugin( + BlockGroupPragma, + new_ops=(BgRoot, BgScope, BgScheduler), + passes=(MaterializeBlockGroup(),), + produces_ops=(BgRoot, BgScope, BgScheduler), + consumes_ops=(MirExecutePipeline,), + ) + class BlockGroupPlugin: + """Doc body lives here. The decorator throws the class away + and returns a PragmaPlugin instance bound to its name.""" + + The decorated object's class body is discarded — only its name is + used (for diagnostics). For richer subclassing (where the plugin + has methods or pre-populated typed fields), construct via + `@final @dataclass(frozen=True, slots=True)` subclassing of + `PragmaPlugin` directly, per the spec § 3.2.1.3 example. + + The bare-constructor form (`PragmaPlugin(pragma_cls=...)`) MUST + also work; this decorator is purely optional sugar. + ''' + + def _wrap(_cls: type) -> PragmaPlugin: + return PragmaPlugin( + pragma_cls=pragma_cls, + new_ops=new_ops, + new_attributes=new_attributes, + passes=passes, + lowerings=lowerings, + renders=renders if renders is not None else {}, + requires_services=requires_services, + produces_ops=produces_ops, + consumes_ops=consumes_ops, + preserves=preserves, + ) + + return _wrap + + +class PragmaPluginConflictError(ValueError): + '''Raised by `Compiler.register_pragma_plugin(plugin)` when the + plugin's contributions conflict with already-registered state. + + Subclasses of this error name the specific conflict kind for + programmatic recovery; the base class is what callers normally + `except` against. + + Per spec § 3.2.1.3: registration is atomic. On conflict, the + per-Compiler registries are left UNCHANGED — no half-loaded plugin + state. See `dialect.py:register_pragma_plugin` for the + try-everything-then-commit transaction shape. + ''' + + +@final +class OpNameCollisionError(PragmaPluginConflictError): + '''A new op in `plugin.new_ops` has the same `__name__` as an op + already registered with the Compiler.''' + + +@final +class AttributeNameCollisionError(PragmaPluginConflictError): + '''A new attribute in `plugin.new_attributes` has the same + `__name__` as an attribute already registered with the Compiler.''' + + +@final +class RenderDoubleRegistrationError(PragmaPluginConflictError): + '''A `(op_type, target)` pair in `plugin.renders` already has a + render registered with the Compiler.''' + + +@final +class MissingRequiredServiceError(PragmaPluginConflictError): + '''A type in `plugin.requires_services` is not registered with the + Compiler's `Services` handle.''' + + +__all__ = [ + 'AttributeNameCollisionError', + 'MissingRequiredServiceError', + 'OpNameCollisionError', + 'PragmaPlugin', + 'PragmaPluginConflictError', + 'Render', + 'RenderDoubleRegistrationError', + 'pragma_plugin', +] diff --git a/src/srdatalog/ir/core/services.py b/src/srdatalog/ir/core/services.py new file mode 100644 index 0000000..5f8cd54 --- /dev/null +++ b/src/srdatalog/ir/core/services.py @@ -0,0 +1,149 @@ +'''Typed-key Services dict — Risk 1 fix from spec § 3.2.1.2. + +`Services` is a per-`Compiler.run` typed-key dict of injected services. +Adding a service = new TYPE, never a new attribute on this class. + +The shape is the MLIR/LLVM AnalysisManager / React ContextProvider +pattern (spec § 11.2 / § 11.4): services are injected, not accumulated +into a fixed-schema bag. A fresh-name generator, a view-layout +provider, a per-`Compiler` plugin registry — each is registered under +its own TYPE key. Type the public methods so `services.get(NameGen)` +infers `NameGen` at the call site. + +Per spec § 3.2.1.2: + + Risk 1: `Services` becomes a named-field carrier. + services.name_gen, services.compiler, ... — each new + service edits the `Services` class. Same fixed-schema + anti-pattern. + Fix: `Services` is a typed-key dict: `services.get[T](T) → T`. + Adding a service = new type, no edit. + +Usage: + + from srdatalog.ir.core.lower_ctx import NameGen + services = Services() + services.register(NameGen, NameGen()) + + ng = services.get(NameGen) # mypy infers NameGen + maybe_ng = services.try_get(NameGen) # mypy infers NameGen | None + +This module is part of `core/`; per D6 it imports no concrete dialect +and no concrete service type. The Services class is the typed shell; +each service is its own type, owned by whoever ships it. +''' + +from __future__ import annotations + +from typing import Any, TypeVar + +# TypeVar over service classes. PEP 695 syntax (`def get[T](...)`) is +# Python 3.12+; this project's `requires-python = ">=3.10"` mandates +# the classic TypeVar form for now. The public surface still types +# correctly under mypy --strict for callers. +_T = TypeVar('_T') + + +class ServiceMissingError(KeyError): + '''Raised by `Services.get(T)` when no instance of T is registered. + + Subclass of `KeyError` so existing `except KeyError` clauses keep + catching it; carries the requested type for diagnostics. + ''' + + def __init__(self, service_type: type) -> None: + self.service_type = service_type + super().__init__( + f'no service of type {service_type.__name__!r} registered; ' + f'register one via `services.register({service_type.__name__}, ' + f')` before this call. See spec § 3.2.1.2 Risk 1.' + ) + + +class ServiceConflictError(ValueError): + '''Raised by `Services.register(T, instance)` when T is already + registered. + + Re-registering would silently overwrite the prior instance — a + surprise for downstream passes that already captured a handle. The + conflict is loud at registration time; if a re-register is + intentional, the caller must explicitly `services.replace(T, ...)` + (NOT yet implemented; YAGNI until a real caller needs it). + ''' + + def __init__(self, service_type: type) -> None: + self.service_type = service_type + super().__init__( + f'service of type {service_type.__name__!r} is already registered; ' + f'register() does not allow silent overwrite. See spec § 3.2.1.2 Risk 1.' + ) + + +class Services: + '''Per-`Compiler.run` typed-key dict of injected services. + + No named-field attributes. Internally a `dict[type, Any]`. The + public methods are typed so mypy infers `services.get(NameGen)` -> + `NameGen` at the call site, NOT `Any`. + + Spec: docs/phase_decomposition_redesign.md § 3.2.1.2 (Risk 1 fix) + and § 3.2.1.3 (consumed by `PragmaPlugin.requires_services` for + conflict detection at plugin-registration time). + + Discipline: + - NEVER add a typed attribute to this class. The whole point is + that the API surface is fixed; new services live in their own + types (Risk 1). + - The internal storage is `_by_type: dict[type, Any]`. The + method-level type vars keep callers type-safe without leaking + the storage shape. + ''' + + __slots__ = ('_by_type',) + + def __init__(self) -> None: + self._by_type: dict[type, Any] = {} + + def register(self, service_type: type[_T], instance: _T) -> None: + '''Register `instance` under the key `service_type`. + + Raises `ServiceConflictError` if `service_type` is already + registered. Re-registration is loud by design — see the + `ServiceConflictError` docstring. + ''' + if service_type in self._by_type: + raise ServiceConflictError(service_type) + self._by_type[service_type] = instance + + def get(self, service_type: type[_T]) -> _T: + '''Look up the instance registered under `service_type`. + + Raises `ServiceMissingError` if no instance is registered. + Callers that prefer a None fallback use `try_get(...)` instead. + ''' + try: + return self._by_type[service_type] # type: ignore[no-any-return] + except KeyError: + raise ServiceMissingError(service_type) from None + + def try_get(self, service_type: type[_T]) -> _T | None: + '''Look up the instance registered under `service_type`, or None. + + Symmetric with `dict.get(key, None)`; never raises. + ''' + return self._by_type.get(service_type) + + def __contains__(self, service_type: object) -> bool: + '''Membership test: `MyService in services`. The argument is + typed as `object` (not `type[_T]`) so callers can write + `if MyService in services:` without mypy complaining about + generic-parameter inference. + ''' + return service_type in self._by_type + + +__all__ = [ + 'ServiceConflictError', + 'ServiceMissingError', + 'Services', +] diff --git a/src/srdatalog/ir/core/strategy.py b/src/srdatalog/ir/core/strategy.py index 6e3fbe1..4ca35bf 100644 --- a/src/srdatalog/ir/core/strategy.py +++ b/src/srdatalog/ir/core/strategy.py @@ -33,7 +33,7 @@ import dataclasses from collections.abc import Callable -from typing import TypeVar +from typing import Any, TypeVar from srdatalog.ir.core.ops import Op @@ -202,11 +202,19 @@ def _map_children(op: Op, transform: Strategy, *, all_or_nothing: bool) -> Op | When `all_or_nothing=False`, transform-fails are treated as identity (Stratego `try_(all)` semantics, used by traversal walkers). + + The framework-owned `attributes` field on the `Op` base (PR-P0, + spec § 3.2.1.2 Risks 2 + 5) is intentionally skipped — it carries + per-pass metadata, NOT IR children. `dataclasses.replace` preserves + the field's value when it is not in the kwargs, so the resulting + op still carries the same `AttributeDict` reference. ''' # Op subclasses are always dataclasses by D4 (see design_principles.md). changed = False - new_values: dict[str, object] = {} + new_values: dict[str, Any] = {} for f in dataclasses.fields(op): + if f.name == 'attributes': + continue val = getattr(op, f.name) new_val, child_changed, ok = _transform_field(val, transform, all_or_nothing) if not ok: @@ -221,14 +229,20 @@ def _map_children(op: Op, transform: Strategy, *, all_or_nothing: bool) -> Op | def _map_one_child(op: Op, transform: Strategy) -> Op | None: '''Apply `transform` to exactly one Op-valued child (left-to-right first that succeeds). Returns None if no child succeeds. + + The framework-owned `attributes` field (PR-P0) is skipped — see + `_map_children` docstring for the rationale. ''' # Op subclasses are always dataclasses by D4 (see design_principles.md). for f in dataclasses.fields(op): + if f.name == 'attributes': + continue val = getattr(op, f.name) if isinstance(val, Op): r = transform(val) if r is not None and r is not val: - return dataclasses.replace(op, **{f.name: r}) + kw: dict[str, Any] = {f.name: r} + return dataclasses.replace(op, **kw) elif isinstance(val, (list, tuple)): for i, x in enumerate(val): if isinstance(x, Op): @@ -236,18 +250,24 @@ def _map_one_child(op: Op, transform: Strategy) -> Op | None: if r is not None and r is not x: new_list = list(val) new_list[i] = r - return dataclasses.replace(op, **{f.name: type(val)(new_list)}) + kw_seq: dict[str, Any] = {f.name: type(val)(new_list)} + return dataclasses.replace(op, **kw_seq) return None def _map_some_children(op: Op, transform: Strategy) -> Op | None: '''Apply `transform` to as many Op-valued children as succeed. Returns None if no child succeeds (Stratego `some` semantics). + + The framework-owned `attributes` field (PR-P0) is skipped — see + `_map_children` docstring for the rationale. ''' # Op subclasses are always dataclasses by D4 (see design_principles.md). any_succeeded = False - new_values: dict[str, object] = {} + new_values: dict[str, Any] = {} for f in dataclasses.fields(op): + if f.name == 'attributes': + continue val = getattr(op, f.name) new_val, ok = _transform_field_some(val, transform) new_values[f.name] = new_val diff --git a/tests/test_pragma_plugin_framework.py b/tests/test_pragma_plugin_framework.py new file mode 100644 index 0000000..565b3ef --- /dev/null +++ b/tests/test_pragma_plugin_framework.py @@ -0,0 +1,856 @@ +'''PR-P0 tests — pragma plugin framework primitives. + +Covers the seven new primitives delivered by PR-P0 +(`docs/phase_decomposition_redesign.md` § 6.0.0): + + 1. `PragmaPlugin` construction (positional + decorator factory). + 2. `Services.register / get / try_get` (typed retrieval). + 3. `AttributeDict.__getitem__ / __setitem__ / __contains__ / get`. + 4. `Compiler.register_pragma_plugin` atomicity (build a plugin + with a conflict; assert no partial state). + 5. `topo_sort_passes` cycle detection. + 6. `_compiler_registration_scope` thread-locality (decorators + stage into compiler when scope active). + 7. Back-compat: legacy `@pragma_handler` still works (minimal + pragma + handler + run through `MaterializePragmaPass`). + +These tests live OUTSIDE `tests/test_core_pragma.py` because the +existing file pins the back-compat surface; PR-P0 introduces new +primitives without disturbing it. +''' + +from __future__ import annotations + +from dataclasses import dataclass +from typing import final + +import pytest + +from srdatalog.ir.core import ( + Attribute, + AttributeDict, + AttributeNameCollisionError, + Compiler, + MaterializePragmaPass, + MissingRequiredServiceError, + Op, + OpNameCollisionError, + Pass, + PassCycleError, + Pragma, + PragmaPlugin, + Render, + RenderDoubleRegistrationError, + ServiceConflictError, + ServiceMissingError, + Services, + pragma_handler, + pragma_plugin, + topo_sort_passes, +) +from srdatalog.ir.core.dialect import _compiler_registration_scope, _get_current_compiler +from srdatalog.ir.core.pragma import PragmaRegistration + +# ============================================================================= +# Test fixtures +# ============================================================================= + + +@pytest.fixture +def isolated_pragma_registry(monkeypatch): + '''Swap the module-global `_PRAGMA_REGISTRY` for an empty list so + each test's decorations stay isolated. Mirrors the existing + `isolated_registry` fixture in `tests/test_core_pragma.py`. + ''' + fresh: list[PragmaRegistration] = [] + monkeypatch.setattr('srdatalog.ir.core.pragma._PRAGMA_REGISTRY', fresh) + return fresh + + +# A toy Pragma + ops + attributes used across multiple tests. + + +@final +@dataclass(frozen=True, slots=True) +class _ToyPragma(Pragma): + capacity: int = 10 + + +@final +@dataclass(frozen=True, slots=True) +class _ToyOp(Op): + name: str = 'toy' + + +@final +@dataclass(frozen=True, slots=True) +class _OtherToyOp(Op): + value: int = 0 + + +@final +@dataclass(frozen=True, slots=True) +class _ToyAttr(Attribute): + threshold: float = 0.5 + + +# A bare service type used to exercise the Services dict. + + +class _NameGenLike: + def __init__(self) -> None: + self.counter = 0 + + def fresh(self, prefix: str) -> str: + self.counter += 1 + return f'{prefix}_{self.counter}' + + +class _ViewLayoutLike: + pass + + +# ============================================================================= +# 1. PragmaPlugin construction +# ============================================================================= + + +def test_pragma_plugin_constructs_with_defaults(): + '''Bare `PragmaPlugin(pragma_cls=...)` works — every other field + defaults to () / empty dict per spec § 3.2.1.3.''' + plugin = PragmaPlugin(pragma_cls=_ToyPragma) + assert plugin.pragma_cls is _ToyPragma + assert plugin.new_ops == () + assert plugin.new_attributes == () + assert plugin.passes == () + assert plugin.lowerings == () + assert plugin.renders == {} + assert plugin.requires_services == () + assert plugin.produces_ops == () + assert plugin.consumes_ops == () + assert plugin.preserves == () + + +def test_pragma_plugin_constructs_with_fields_populated(): + '''Every field on `PragmaPlugin` is settable via the constructor + (the spec § 3.2.1.3 example uses all of them).''' + plugin = PragmaPlugin( + pragma_cls=_ToyPragma, + new_ops=(_ToyOp,), + new_attributes=(_ToyAttr,), + requires_services=(_NameGenLike,), + produces_ops=(_ToyOp,), + consumes_ops=(_OtherToyOp,), + preserves=(_NameGenLike,), + ) + assert plugin.new_ops == (_ToyOp,) + assert plugin.new_attributes == (_ToyAttr,) + assert plugin.requires_services == (_NameGenLike,) + assert plugin.produces_ops == (_ToyOp,) + assert plugin.consumes_ops == (_OtherToyOp,) + assert plugin.preserves == (_NameGenLike,) + + +def test_pragma_plugin_is_frozen(): + '''`PragmaPlugin` is a frozen dataclass; field reassignment fails.''' + from dataclasses import FrozenInstanceError + + plugin = PragmaPlugin(pragma_cls=_ToyPragma) + with pytest.raises(FrozenInstanceError): + plugin.pragma_cls = _ToyPragma # type: ignore[misc] + + +def test_pragma_plugin_decorator_factory_builds_a_plugin(): + '''`@pragma_plugin(...)` class-decorator factory returns a + PragmaPlugin instance bound to the class name.''' + + @pragma_plugin( + _ToyPragma, + new_ops=(_ToyOp,), + produces_ops=(_ToyOp,), + ) + class _MyToyPlugin: + """Sugar form: class body discarded, decorator returns the plugin.""" + + assert isinstance(_MyToyPlugin, PragmaPlugin) + assert _MyToyPlugin.pragma_cls is _ToyPragma + assert _MyToyPlugin.new_ops == (_ToyOp,) + assert _MyToyPlugin.produces_ops == (_ToyOp,) + + +def test_pragma_plugin_decorator_renders_default_is_empty_dict(): + '''The decorator factory's `renders=None` default becomes `{}` on + the resulting plugin (not None — avoids the mutable-default trap + for the constructor's `default_factory=dict`).''' + + @pragma_plugin(_ToyPragma) + class _Plugin: + """No renders contributed.""" + + assert isinstance(_Plugin, PragmaPlugin) + assert _Plugin.renders == {} + + +# ============================================================================= +# 2. Services typed-key dict +# ============================================================================= + + +def test_services_register_and_get(): + '''Round-trip: register an instance, look it up via its type.''' + services = Services() + ng = _NameGenLike() + services.register(_NameGenLike, ng) + assert services.get(_NameGenLike) is ng + + +def test_services_get_returns_the_exact_instance(): + '''No deep copy — the same object is returned.''' + services = Services() + vl = _ViewLayoutLike() + services.register(_ViewLayoutLike, vl) + assert services.get(_ViewLayoutLike) is vl + + +def test_services_get_raises_when_missing(): + '''Looking up an unregistered type raises ServiceMissingError + (subclass of KeyError so legacy `except KeyError` keeps catching).''' + services = Services() + with pytest.raises(ServiceMissingError) as ei: + services.get(_NameGenLike) + assert issubclass(type(ei.value), KeyError) + assert ei.value.service_type is _NameGenLike + + +def test_services_try_get_returns_none_when_missing(): + '''try_get is the nullable form; never raises.''' + services = Services() + assert services.try_get(_NameGenLike) is None + + +def test_services_try_get_returns_instance_when_present(): + services = Services() + ng = _NameGenLike() + services.register(_NameGenLike, ng) + assert services.try_get(_NameGenLike) is ng + + +def test_services_register_rejects_double_registration(): + '''Re-registering the same TYPE raises ServiceConflictError — + silent overwrites would mismatch downstream captured handles.''' + services = Services() + services.register(_NameGenLike, _NameGenLike()) + with pytest.raises(ServiceConflictError) as ei: + services.register(_NameGenLike, _NameGenLike()) + assert ei.value.service_type is _NameGenLike + + +def test_services_contains(): + '''`MyService in services` returns True iff registered.''' + services = Services() + assert _NameGenLike not in services + services.register(_NameGenLike, _NameGenLike()) + assert _NameGenLike in services + + +def test_services_typed_get_infers_concrete_type(): + '''mypy/runtime parity: get(NameGenLike) returns a NameGenLike, + not Any. The runtime check here pairs with the typing tests in + the module's stub: `services.get(_NameGenLike).fresh(...)` + type-checks at mypy and works at runtime.''' + services = Services() + services.register(_NameGenLike, _NameGenLike()) + ng = services.get(_NameGenLike) + # Runtime: the returned value supports the protocol of _NameGenLike + assert ng.fresh('x') == 'x_1' + + +# ============================================================================= +# 3. AttributeDict typed-key dict +# ============================================================================= + + +def test_attribute_dict_set_and_get_via_indexing(): + attrs = AttributeDict() + v = _ToyAttr(threshold=0.7) + attrs[_ToyAttr] = v + assert attrs[_ToyAttr] is v + assert attrs[_ToyAttr].threshold == 0.7 + + +def test_attribute_dict_getitem_raises_when_missing(): + attrs = AttributeDict() + with pytest.raises(KeyError) as ei: + _ = attrs[_ToyAttr] + assert '_ToyAttr' in str(ei.value) + + +def test_attribute_dict_get_returns_none_when_missing(): + attrs = AttributeDict() + assert attrs.get(_ToyAttr) is None + + +def test_attribute_dict_get_returns_value_when_present(): + attrs = AttributeDict() + v = _ToyAttr() + attrs[_ToyAttr] = v + assert attrs.get(_ToyAttr) is v + + +def test_attribute_dict_contains(): + attrs = AttributeDict() + assert _ToyAttr not in attrs + attrs[_ToyAttr] = _ToyAttr() + assert _ToyAttr in attrs + + +def test_attribute_dict_setitem_validates_value_type(): + '''Writing a mismatched value type fails at write time, not at + the eventual read crash.''' + + @final + @dataclass(frozen=True, slots=True) + class OtherAttr(Attribute): + pass + + attrs = AttributeDict() + with pytest.raises(TypeError, match=r'requires an instance of'): + attrs[_ToyAttr] = OtherAttr() # type: ignore[assignment] + + +def test_attribute_dict_setitem_overwrites_silently(): + '''Repeat writes of the SAME type silently overwrite — per the + class docstring, conflict detection is the producer's responsibility.''' + attrs = AttributeDict() + v1 = _ToyAttr(threshold=0.1) + v2 = _ToyAttr(threshold=0.9) + attrs[_ToyAttr] = v1 + attrs[_ToyAttr] = v2 + assert attrs[_ToyAttr] is v2 + + +def test_attribute_dict_len_and_iter(): + '''Mirror dict semantics: len() and iter() over keys.''' + attrs = AttributeDict() + assert len(attrs) == 0 + attrs[_ToyAttr] = _ToyAttr() + assert len(attrs) == 1 + assert list(attrs) == [_ToyAttr] + + +def test_attribute_dict_repr(): + attrs = AttributeDict() + assert repr(attrs) == 'AttributeDict()' + attrs[_ToyAttr] = _ToyAttr(threshold=0.3) + assert '_ToyAttr' in repr(attrs) + + +def test_op_default_attributes_is_a_fresh_dict_per_instance(): + '''Two ops with default-factory `attributes` get independent dicts + (the `field(default_factory=...)` is per-instance).''' + + @dataclass(frozen=True, slots=True) + class MyOp(Op): + name: str = 'a' + + a, b = MyOp(), MyOp() + assert a.attributes is not b.attributes + a.attributes[_ToyAttr] = _ToyAttr() + assert _ToyAttr not in b.attributes + + +def test_op_attributes_does_not_affect_equality(): + '''Two ops with the same fields but different attribute states + are still equal — `compare=False` on the parent's `attributes` + field excludes it from structural identity.''' + + @dataclass(frozen=True, slots=True) + class MyOp(Op): + name: str = 'a' + + a, b = MyOp(), MyOp() + a.attributes[_ToyAttr] = _ToyAttr() + assert a == b + assert hash(a) == hash(b) + + +def test_op_attributes_excluded_from_repr(): + '''`repr=False` on the parent's `attributes` field keeps repr + byte-equivalent with the pre-PR-P0 shape.''' + + @dataclass(frozen=True, slots=True) + class MyOp(Op): + name: str = 'a' + + r = repr(MyOp(name='foo')) + # The repr should NOT include the attributes= kwarg or its + # AttributeDict marker. + assert 'attributes=' not in r + assert 'AttributeDict' not in r + + +# ============================================================================= +# 4. Compiler.register_pragma_plugin atomicity +# ============================================================================= + + +def _make_simple_plugin(pragma_cls: type[Pragma] = _ToyPragma) -> PragmaPlugin: + '''A baseline plugin contributing one new op + one new attribute.''' + return PragmaPlugin( + pragma_cls=pragma_cls, + new_ops=(_ToyOp,), + new_attributes=(_ToyAttr,), + ) + + +def test_register_pragma_plugin_succeeds_on_clean_compiler(): + c = Compiler() + plugin = _make_simple_plugin() + c.register_pragma_plugin(plugin) + assert '_ToyOp' in c.ops + assert c.ops['_ToyOp'] is _ToyOp + assert '_ToyAttr' in c.attributes + assert plugin in c.pragma_plugins + + +def test_register_pragma_plugin_op_name_collision_is_atomic(): + '''Op-name collision aborts the whole registration; the second + plugin's other contributions are NOT applied.''' + + # First plugin + c = Compiler() + c.register_pragma_plugin(PragmaPlugin(pragma_cls=_ToyPragma, new_ops=(_ToyOp,))) + + # Second plugin with a different Op class but same __name__. + # Build a synthetic class with the same name. + @dataclass(frozen=True, slots=True) + class _Shadow(Op): + pass + + _Shadow.__name__ = '_ToyOp' + + attrs_before = dict(c.attributes) + ops_before = dict(c.ops) + passes_before = list(c.passes) + with pytest.raises(OpNameCollisionError): + c.register_pragma_plugin( + PragmaPlugin( + pragma_cls=_ToyPragma, + new_ops=(_Shadow,), + new_attributes=(_ToyAttr,), + ) + ) + # No state changed: attributes / passes untouched + assert c.ops == ops_before + assert c.attributes == attrs_before + assert c.passes == passes_before + + +def test_register_pragma_plugin_attribute_name_collision_is_atomic(): + c = Compiler() + c.register_pragma_plugin(PragmaPlugin(pragma_cls=_ToyPragma, new_attributes=(_ToyAttr,))) + + @dataclass(frozen=True, slots=True) + class _ShadowAttr(Attribute): + pass + + _ShadowAttr.__name__ = '_ToyAttr' + + ops_before = dict(c.ops) + with pytest.raises(AttributeNameCollisionError): + c.register_pragma_plugin( + PragmaPlugin( + pragma_cls=_ToyPragma, + new_ops=(_ToyOp,), + new_attributes=(_ShadowAttr,), + ) + ) + # The op contribution from the failing plugin was NOT applied + assert '_ToyOp' not in ops_before + assert '_ToyOp' not in c.ops + + +def test_register_pragma_plugin_render_double_registration_is_atomic(): + '''Two plugins registering a render for the same (op_type, target) + pair: second registration fails atomically.''' + c = Compiler() + + def _render_a(op, ctx): + return 'a' + + def _render_b(op, ctx): + return 'b' + + c.register_pragma_plugin( + PragmaPlugin( + pragma_cls=_ToyPragma, + renders={'cuda': (Render(op_type=_ToyOp, fn=_render_a),)}, + ) + ) + + # Second plugin tries to render the same (op_type, target) pair. + ops_before = dict(c.ops) + with pytest.raises(RenderDoubleRegistrationError): + c.register_pragma_plugin( + PragmaPlugin( + pragma_cls=_ToyPragma, + new_ops=(_OtherToyOp,), + renders={'cuda': (Render(op_type=_ToyOp, fn=_render_b),)}, + ) + ) + # `_OtherToyOp` contribution was rolled back + assert '_OtherToyOp' not in c.ops + assert dict(c.ops) == ops_before + + +def test_register_pragma_plugin_missing_required_service_is_atomic(): + c = Compiler() + ops_before = dict(c.ops) + with pytest.raises(MissingRequiredServiceError) as ei: + c.register_pragma_plugin( + PragmaPlugin( + pragma_cls=_ToyPragma, + new_ops=(_ToyOp,), + requires_services=(_NameGenLike,), + ) + ) + assert '_NameGenLike' in str(ei.value) + # No op contribution from the failing plugin + assert c.ops == ops_before + + +def test_register_pragma_plugin_satisfied_service_succeeds(): + c = Compiler() + c.services.register(_NameGenLike, _NameGenLike()) + c.register_pragma_plugin( + PragmaPlugin( + pragma_cls=_ToyPragma, + new_ops=(_ToyOp,), + requires_services=(_NameGenLike,), + ) + ) + assert '_ToyOp' in c.ops + + +def test_register_pragma_plugin_pass_cycle_is_atomic(): + '''Two passes whose produces/consumes form a cycle: registration + fails, neither pass is added to `compiler.passes`.''' + + @dataclass(frozen=True) + class _PA(Pass): + name: str = 'pa' + produces_ops: tuple[type, ...] = (_ToyOp,) + consumes_ops: tuple[type, ...] = (_OtherToyOp,) + + def apply(self, prog, compiler): + return prog + + @dataclass(frozen=True) + class _PB(Pass): + name: str = 'pb' + produces_ops: tuple[type, ...] = (_OtherToyOp,) + consumes_ops: tuple[type, ...] = (_ToyOp,) + + def apply(self, prog, compiler): + return prog + + c = Compiler() + passes_before = list(c.passes) + with pytest.raises(PassCycleError): + c.register_pragma_plugin( + PragmaPlugin( + pragma_cls=_ToyPragma, + passes=(_PA(), _PB()), + ) + ) + # No partial pass state + assert c.passes == passes_before + + +# ============================================================================= +# 5. topo_sort_passes cycle detection +# ============================================================================= + + +def test_topo_sort_passes_returns_input_when_no_deps(): + '''An empty list of passes returns an empty list. Independent + passes preserve a stable order keyed on name.''' + + @dataclass(frozen=True) + class _P(Pass): + name: str = 'p' + + def apply(self, prog, compiler): + return prog + + assert topo_sort_passes([]) == [] + out = topo_sort_passes([_P(name='b'), _P(name='a')]) + assert [p.name for p in out] == ['a', 'b'] + + +def test_topo_sort_passes_respects_produces_consumes(): + '''A pass that produces an op must run before any pass that + consumes it.''' + + @dataclass(frozen=True) + class _Producer(Pass): + name: str = 'producer' + produces_ops: tuple[type, ...] = (_ToyOp,) + + def apply(self, prog, compiler): + return prog + + @dataclass(frozen=True) + class _Consumer(Pass): + name: str = 'consumer' + consumes_ops: tuple[type, ...] = (_ToyOp,) + + def apply(self, prog, compiler): + return prog + + out = topo_sort_passes([_Consumer(), _Producer()]) + assert [p.name for p in out] == ['producer', 'consumer'] + + +def test_topo_sort_passes_self_edge_tolerated(): + '''A pass that both produces and consumes the same op type is + allowed (the materialization passes do this); no cycle is raised.''' + + @dataclass(frozen=True) + class _Materializer(Pass): + name: str = 'mat' + produces_ops: tuple[type, ...] = (_ToyOp,) + consumes_ops: tuple[type, ...] = (_ToyOp,) + + def apply(self, prog, compiler): + return prog + + out = topo_sort_passes([_Materializer()]) + assert len(out) == 1 + + +def test_topo_sort_passes_raises_on_cycle(): + '''An A->B->A cycle raises PassCycleError naming both passes.''' + + @dataclass(frozen=True) + class _A(Pass): + name: str = 'a' + produces_ops: tuple[type, ...] = (_ToyOp,) + consumes_ops: tuple[type, ...] = (_OtherToyOp,) + + def apply(self, prog, compiler): + return prog + + @dataclass(frozen=True) + class _B(Pass): + name: str = 'b' + produces_ops: tuple[type, ...] = (_OtherToyOp,) + consumes_ops: tuple[type, ...] = (_ToyOp,) + + def apply(self, prog, compiler): + return prog + + with pytest.raises(PassCycleError) as ei: + topo_sort_passes([_A(), _B()]) + assert set(ei.value.cycle) == {'a', 'b'} + + +# ============================================================================= +# 6. _compiler_registration_scope thread-locality +# ============================================================================= + + +def test_compiler_registration_scope_sets_thread_local(): + c = Compiler() + assert _get_current_compiler() is None + with _compiler_registration_scope(c): + assert _get_current_compiler() is c + assert _get_current_compiler() is None + + +def test_compiler_registration_scope_nests_and_restores(): + c1 = Compiler() + c2 = Compiler() + with _compiler_registration_scope(c1): + assert _get_current_compiler() is c1 + with _compiler_registration_scope(c2): + assert _get_current_compiler() is c2 + assert _get_current_compiler() is c1 + assert _get_current_compiler() is None + + +def test_compiler_registration_scope_restores_on_exception(): + c = Compiler() + with pytest.raises(RuntimeError), _compiler_registration_scope(c): + assert _get_current_compiler() is c + raise RuntimeError('boom') + assert _get_current_compiler() is None + + +def test_pragma_handler_stages_into_compiler_when_scope_active(isolated_pragma_registry): + '''When a `_compiler_registration_scope(c)` is active, a + `@pragma_handler` decoration stages its registration into the + compiler's per-instance handler list — NOT the module-global.''' + + c = Compiler() + + with _compiler_registration_scope(c): + + @pragma_handler(_ToyPragma, on=_ToyOp) + def _handler(op, pragma, ctx): + return op + + # Module-global registry stayed empty + assert isolated_pragma_registry == [] + # Per-Compiler list got the registration + staged: list[PragmaRegistration] = getattr(c, '_pragma_handlers', []) + assert len(staged) == 1 + assert staged[0].pragma_cls is _ToyPragma + + +def test_pragma_handler_falls_back_to_module_global_without_scope(isolated_pragma_registry): + '''Outside a registration scope, `@pragma_handler` keeps writing + to the module-global (back-compat with pre-PR-P0 code paths).''' + + @pragma_handler(_ToyPragma, on=_ToyOp) + def _handler(op, pragma, ctx): + return op + + assert len(isolated_pragma_registry) == 1 + assert isolated_pragma_registry[0].pragma_cls is _ToyPragma + + +# ============================================================================= +# 7. Back-compat: legacy @pragma_handler + MaterializePragmaPass + Compiler.run +# ============================================================================= + + +def test_materialize_pragma_pass_runs_handler_via_compiler_run(isolated_pragma_registry): + '''End-to-end back-compat: a `@pragma_handler` registered inside a + compiler scope, then a `MaterializePragmaPass(_ToyPragma)` Pass + inserted into `compiler.passes`, then `compiler.run(prog)` invokes + the handler against matching ops in `prog`. + + This validates spec § 6.0.0 row 8 (the legacy `@pragma_handler` + back-compat shim). + ''' + + # An op carrying pragmas (mimics mir.ExecutePipeline shape). + @final + @dataclass(frozen=True, slots=True) + class _Carrier(Op): + name: str = '' + pragmas: tuple[Pragma, ...] = () + + # Handler returns a NEW _Carrier with `name='materialized'`. + call_log: list[str] = [] + + c = Compiler() + with _compiler_registration_scope(c): + + @pragma_handler(_ToyPragma, on=_Carrier) + def _materialize(op, pragma, ctx): + call_log.append(f'fired:{op.name}') + return _Carrier(name='materialized', pragmas=()) + + # Build the pipeline using the MaterializePragmaPass shim. + c.passes.append(MaterializePragmaPass(name='mat', pragma_cls=_ToyPragma)) + + # Run against a Carrier op carrying the toy pragma. + prog = _Carrier(name='before', pragmas=(_ToyPragma(),)) + result = c.run(prog) + + assert call_log == ['fired:before'] + assert isinstance(result, _Carrier) + assert result.name == 'materialized' + + +def test_materialize_pragma_pass_is_noop_when_no_pragma_matches(isolated_pragma_registry): + '''Carrier op without the matching pragma → handler never fires; + op returned unchanged.''' + + @final + @dataclass(frozen=True, slots=True) + class _Carrier(Op): + pragmas: tuple[Pragma, ...] = () + + call_log: list[str] = [] + + c = Compiler() + with _compiler_registration_scope(c): + + @pragma_handler(_ToyPragma, on=_Carrier) + def _materialize(op, pragma, ctx): + call_log.append('fired') + return op + + c.passes.append(MaterializePragmaPass(name='mat', pragma_cls=_ToyPragma)) + prog = _Carrier(pragmas=()) + result = c.run(prog) + assert call_log == [] + assert result is prog + + +def test_compiler_run_uses_self_passes_when_no_pipeline_given(): + '''PR-P0 (spec § 6 row 6): `compiler.run(prog)` without + `pipeline=...` consults `self.passes` directly. Existing + `compiler.run(prog, pipeline=[...])` continues to work + (back-compat). + ''' + + @dataclass(frozen=True) + class _Identity(Pass): + name: str = 'identity' + + def apply(self, prog, compiler): + return prog + + c = Compiler() + c.passes.append(_Identity()) + + out = c.run('hello') + assert out == 'hello' + + out2 = c.run('explicit', pipeline=[_Identity()]) + assert out2 == 'explicit' + + +# ============================================================================= +# 8. Per-Compiler registries are independent across instances +# ============================================================================= + + +def test_per_compiler_registries_do_not_leak_across_instances(): + '''Risk-4 fix: registries are per-Compiler. A plugin registered on + c1 does NOT show up in c2.''' + c1 = Compiler() + c2 = Compiler() + c1.register_pragma_plugin(PragmaPlugin(pragma_cls=_ToyPragma, new_ops=(_ToyOp,))) + assert '_ToyOp' in c1.ops + assert '_ToyOp' not in c2.ops + + +# ============================================================================= +# 9. Sanity: PragmaPlugin __all__ surface +# ============================================================================= + + +def test_pragma_plugin_imports_are_public(): + '''Every NEW PR-P0 primitive must be importable from the public + package surface (`srdatalog.ir.core.__all__`).''' + from srdatalog.ir.core import __all__ as core_all + + expected = { + 'Attribute', + 'AttributeDict', + 'MaterializePragmaPass', + 'PassCycleError', + 'PragmaPlugin', + 'PragmaPluginConflictError', + 'Render', + 'Services', + 'pragma_plugin', + 'topo_sort_passes', + } + missing = expected - set(core_all) + assert not missing, f'Missing from srdatalog.ir.core.__all__: {missing}'