diff --git a/src/srdatalog/ir/core/dialect.py b/src/srdatalog/ir/core/dialect.py index c7a62bf..11f12b3 100644 --- a/src/srdatalog/ir/core/dialect.py +++ b/src/srdatalog/ir/core/dialect.py @@ -117,6 +117,13 @@ def __init__(self) -> None: # to the right plugin. self._active_plugin_context: str | None = None self._active_plugin_replaces: tuple[str, ...] = () + # Active render target for the in-flight `Compiler.run(...)` call. + # Set on entry, restored on exit. Outside a `run` call this is the + # process-wide default. PR-1c (Phase B2-1, foundation) introduces + # this attr so downstream passes can read the target the caller + # chose; the actual per-target dispatch logic lands in Phase B2-2. + # See docs/phase_decomposition_redesign.md §3.3.1. + self.active_target: str = 'cuda' # -- dialect registration -------------------------------------------------- @@ -174,7 +181,7 @@ def dialects(self) -> list[Dialect]: '''All registered dialects, in registration order.''' return list(self._dialects.values()) - def run(self, prog: Any, *, pipeline: list[Any]) -> Any: + def run(self, prog: Any, *, pipeline: list[Any], target: str = 'cuda') -> Any: '''Drive a list of `Pass` instances over `prog`. Returns the (possibly transformed) prog. @@ -184,6 +191,13 @@ def run(self, prog: Any, *, pipeline: list[Any]) -> Any: `PassOrderingError` on mismatch — at construction time, before any pass executes. + `target` selects the render target for the run; passes that need + it read `compiler.active_target` (set for the duration of the + call, restored on exit). Defaults to ``'cuda'`` so all existing + call sites are byte-equivalent. PR-1c (Phase B2-1, foundation) + introduces the kwarg; actual per-target dispatch lands in Phase + B2-2. See docs/phase_decomposition_redesign.md §3.3.1. + Per `docs/compiler_redesign.md` §4 and the R4 research report: pipelines are data; ordering errors are caught up-front. ''' @@ -198,9 +212,14 @@ def run(self, prog: Any, *, pipeline: list[Any]) -> Any: raise PassOrderingError(p.name, needed, i) available |= set(p.produces) - for p in pipeline: - prog = p.apply(prog, self) - return prog + prev_target = self.active_target + self.active_target = target + try: + for p in pipeline: + prog = p.apply(prog, self) + return prog + finally: + self.active_target = prev_target # -- plugin registration --------------------------------------------------- diff --git a/src/srdatalog/ir/default_pipelines.py b/src/srdatalog/ir/default_pipelines.py index 990ec3f..860ecbe 100644 --- a/src/srdatalog/ir/default_pipelines.py +++ b/src/srdatalog/ir/default_pipelines.py @@ -53,13 +53,20 @@ class InitialProg: Read by `HirPlanningShim` when computing HIR from scratch; ignored if `hir` is pre-populated by the caller. Added in F5.2 so the declarative pipeline can preserve the imperative entry point's - verbosity knob byte-equivalently.''' + verbosity knob byte-equivalently. + + `target` carries the render target chosen by the caller. PR-1c + (Phase B2-1, foundation) adds the field so downstream pipeline + stages can read it; default ``'cuda'`` preserves all existing call + sites byte-equivalently. Actual per-target dispatch lands in Phase + B2-2. See docs/phase_decomposition_redesign.md §3.3.1.''' program: Program hir: HirProgram | None = None steps: list[tuple[mir.MirNode, bool]] | None = None mir_program: mir.Program | None = None verbose: bool = False + target: str = 'cuda' @dataclass(frozen=True, slots=True) @@ -67,7 +74,14 @@ class KernelCtx: '''Kernel-pipeline through-state. Carries the original `compile_kernel_body` kwargs plus the progressively populated derived state (handle-assigned ep, view_specs, view_decls + var - maps, iir, body_text).''' + maps, iir, body_text). + + `target` carries the render target chosen by the caller. PR-1c + (Phase B2-1, foundation) adds the field so kernel-pipeline shims + (notably `VerifyRenderabilityShim`) can read it; default ``'cuda'`` + preserves all existing call sites byte-equivalently. Actual + per-target render dispatch lands in Phase B2-2. See + docs/phase_decomposition_redesign.md §3.3.1.''' ep: ExecutePipeline is_counting: bool @@ -77,6 +91,7 @@ class KernelCtx: rel_index_types: dict[str, str] | None = None tiled_cartesian: bool = False bg_enabled: bool = False + target: str = 'cuda' # populated by shims: view_specs: tuple[ViewSpec, ...] | None = None view_decls: str | None = None @@ -304,11 +319,11 @@ class VerifyRenderabilityShim(ProgramPass): Runs `verify_renderability` on the post-fixpoint IIR tree. For every op type reachable in the tree, asserts either a registered - `@register_render(op_type, target='cuda')` OR a `@rewrite` is + `@register_render(op_type, target=state.target)` OR a `@rewrite` is present. On failure, raises `UnrenderableOpError` naming the op - type + target — long before the CUDA renderer's own `KeyError` - would surface, and with a precise message pointing at the missing - plugin contribution. + type + target — long before the renderer's own `KeyError` would + surface, and with a precise message pointing at the missing plugin + contribution. Today's IIR rewrite step is conceptual — `LowerKernelBodyShim` produces the final IIR directly. So this shim runs AFTER @@ -321,7 +336,12 @@ class VerifyRenderabilityShim(ProgramPass): introduce a new pseudo-dialect; it's a gate, not a transformation. Per the spec note in `code_discipline.md` R3, this is exactly the "pipeline-stage closure verification" that prevents silent - fall-through in `RenderShim`.''' + fall-through in `RenderShim`. + + PR-1c (Phase B2-1, foundation) replaced the hardcoded + ``target='cuda'`` with ``state.target`` so the check follows the + caller-chosen render target. See + docs/phase_decomposition_redesign.md §2.3 and §3.3.1.''' def __init__(self) -> None: super().__init__( @@ -336,7 +356,7 @@ def _fn(state: KernelCtx, compiler: Any) -> KernelCtx: from srdatalog.ir.core.verifier import verify_renderability assert state.iir is not None, 'VerifyRenderabilityShim: iir not set' - verify_renderability(state.iir, target='cuda', compiler=compiler) + verify_renderability(state.iir, target=state.target, compiler=compiler) return state diff --git a/tests/test_compiler_target_kwarg.py b/tests/test_compiler_target_kwarg.py new file mode 100644 index 0000000..01afdb6 --- /dev/null +++ b/tests/test_compiler_target_kwarg.py @@ -0,0 +1,208 @@ +'''PR-1c (Phase B2-1, foundation) — target kwarg threading. + +Spec: docs/phase_decomposition_redesign.md §2.3 (target monopolization) +and §3.3.1 (target chosen at run time). PR-1c adds the API surface +without changing render-dispatch behavior; default ``target='cuda'`` +preserves all existing call sites byte-equivalently. + +These tests assert the kwarg + field threading is in place AND that +the default 'cuda' path produces identical output to the pre-PR-1c +code path (no compile or no kwarg). +''' + +from __future__ import annotations + +import srdatalog.ir.mir.types as mir +from srdatalog.dsl import Program, Relation, Var +from srdatalog.ir.core import Compiler +from srdatalog.ir.default_pipelines import ( + DEFAULT_KERNEL_PIPELINE, + DEFAULT_PROGRAM_PIPELINE, + InitialProg, + KernelCtx, +) + +# ----------------------------------------------------------------------------- +# Fixtures +# ----------------------------------------------------------------------------- + + +def _tc_program() -> Program: + '''Tiny TC-style program used across the PR-1c kwarg tests.''' + X, Y, Z = Var('x'), Var('y'), Var('z') + arc = Relation('ArcInput', 2) + edge = Relation('Edge', 2) + path = Relation('Path', 2) + return Program( + rules=[ + (edge(X, Y) <= arc(X, Y)).named('EdgeLoad'), + (path(X, Y) <= edge(X, Y)).named('TCBase'), + (path(X, Z) <= path(X, Y) & edge(Y, Z)).named('TCRec'), + ], + ) + + +def _find_execute_pipelines(node: object) -> list[mir.ExecutePipeline]: + out: list[mir.ExecutePipeline] = [] + if isinstance(node, mir.ExecutePipeline): + out.append(node) + elif isinstance(node, mir.FixpointPlan): + for inst in node.instructions: + out.extend(_find_execute_pipelines(inst)) + elif isinstance(node, mir.ParallelGroup): + for op in node.ops: + out.extend(_find_execute_pipelines(op)) + return out + + +def _first_ep(mir_prog: mir.Program) -> mir.ExecutePipeline: + for step, _is_rec in mir_prog.steps: + eps = _find_execute_pipelines(step) + if eps: + return eps[0] + raise AssertionError('no ExecutePipeline in mir_prog') + + +# ----------------------------------------------------------------------------- +# Carrier dataclass field tests +# ----------------------------------------------------------------------------- + + +def test_initial_prog_default_target_is_cuda(): + '''InitialProg.target defaults to 'cuda' (preserves existing + behavior; per spec §3.3.1).''' + prog = InitialProg(program=_tc_program()) + assert prog.target == 'cuda' + + +def test_initial_prog_target_is_replaceable(): + '''The target field is plumbed through `dataclasses.replace`-style + state threading without rejecting non-cuda values.''' + import dataclasses + + prog = InitialProg(program=_tc_program(), target='cpu_tbb') + assert prog.target == 'cpu_tbb' + prog2 = dataclasses.replace(prog, target='wasm') + assert prog2.target == 'wasm' + + +def test_kernel_ctx_default_target_is_cuda(): + '''KernelCtx.target defaults to 'cuda'.''' + compiler = Compiler.with_default_plugins() + state = compiler.run(InitialProg(program=_tc_program()), pipeline=DEFAULT_PROGRAM_PIPELINE) + assert state.mir_program is not None + ep = _first_ep(state.mir_program) + ctx = KernelCtx(ep=ep, is_counting=False) + assert ctx.target == 'cuda' + + +def test_kernel_ctx_target_explicit_round_trip(): + '''Explicit `target='cuda'` round-trips identically to the default.''' + compiler = Compiler.with_default_plugins() + state = compiler.run(InitialProg(program=_tc_program()), pipeline=DEFAULT_PROGRAM_PIPELINE) + assert state.mir_program is not None + ep = _first_ep(state.mir_program) + + default = compiler.run(KernelCtx(ep=ep, is_counting=False), pipeline=DEFAULT_KERNEL_PIPELINE) + explicit = compiler.run( + KernelCtx(ep=ep, is_counting=False, target='cuda'), pipeline=DEFAULT_KERNEL_PIPELINE + ) + assert default.body_text == explicit.body_text + + +# ----------------------------------------------------------------------------- +# Compiler.run target kwarg +# ----------------------------------------------------------------------------- + + +def test_compiler_run_accepts_target_kwarg(): + '''`Compiler.run(..., target='cuda')` is accepted and stored on the + compiler for the duration of the call.''' + compiler = Compiler.with_default_plugins() + state = compiler.run( + InitialProg(program=_tc_program()), + pipeline=DEFAULT_PROGRAM_PIPELINE, + target='cuda', + ) + assert state.mir_program is not None + + +def test_compiler_run_default_target_is_cuda(): + '''Omitting the kwarg defaults to 'cuda'; the active_target attr is + restored to the prior value after the run.''' + compiler = Compiler.with_default_plugins() + assert compiler.active_target == 'cuda' + + state = compiler.run(InitialProg(program=_tc_program()), pipeline=DEFAULT_PROGRAM_PIPELINE) + assert state.mir_program is not None + # Restored after the call returns. + assert compiler.active_target == 'cuda' + + +def test_compiler_run_target_kwarg_is_byte_equivalent_to_default(): + '''Threading target='cuda' through the kernel pipeline produces + byte-identical body_text to the default (no-kwarg) path. This is + the load-bearing byte-equivalence guarantee of PR-1c: the kwarg is + pure plumbing today.''' + compiler = Compiler.with_default_plugins() + prog_state = compiler.run(InitialProg(program=_tc_program()), pipeline=DEFAULT_PROGRAM_PIPELINE) + assert prog_state.mir_program is not None + ep = _first_ep(prog_state.mir_program) + + default = compiler.run( + KernelCtx(ep=ep, is_counting=False), + pipeline=DEFAULT_KERNEL_PIPELINE, + ) + explicit = compiler.run( + KernelCtx(ep=ep, is_counting=False, target='cuda'), + pipeline=DEFAULT_KERNEL_PIPELINE, + target='cuda', + ) + assert default.body_text == explicit.body_text + assert default.body_text is not None + + +def test_compiler_run_target_restored_on_exception(): + '''If a pass raises, `active_target` must still be restored to the + prior value (try/finally contract).''' + import contextlib + + from srdatalog.ir.core import Pass + + class _Boom(Pass): + def __init__(self) -> None: + super().__init__(name='_boom') + + def apply(self, prog: object, compiler: object) -> object: + raise RuntimeError('boom') + + compiler = Compiler.with_default_plugins() + assert compiler.active_target == 'cuda' + with contextlib.suppress(RuntimeError): + compiler.run(object(), pipeline=[_Boom()], target='cpu_tbb') + assert compiler.active_target == 'cuda' + + +# ----------------------------------------------------------------------------- +# VerifyRenderabilityShim reads target from KernelCtx +# ----------------------------------------------------------------------------- + + +def test_verify_renderability_uses_kernel_ctx_target(): + '''`VerifyRenderabilityShim` reads `state.target` (no longer hardcodes + 'cuda'). With default target='cuda' the existing renderers cover the + IIR tree, so the pipeline runs cleanly — the regression we'd see if + someone re-hardcoded the literal is this test failing for a non-cuda + KernelCtx, but the byte-equivalence path with 'cuda' must stay + green.''' + compiler = Compiler.with_default_plugins() + prog_state = compiler.run(InitialProg(program=_tc_program()), pipeline=DEFAULT_PROGRAM_PIPELINE) + assert prog_state.mir_program is not None + ep = _first_ep(prog_state.mir_program) + + # Default target='cuda' renders cleanly through the verifier. + out = compiler.run( + KernelCtx(ep=ep, is_counting=False, target='cuda'), + pipeline=DEFAULT_KERNEL_PIPELINE, + ) + assert out.body_text is not None