diff --git a/docs/plans/2026-03-21-cython-runtime-core-alignment.md b/docs/plans/2026-03-21-cython-runtime-core-alignment.md new file mode 100644 index 00000000..ec540730 --- /dev/null +++ b/docs/plans/2026-03-21-cython-runtime-core-alignment.md @@ -0,0 +1,811 @@ +# Cython Runtime Core Alignment Implementation Plan + +> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. + +**Goal:** Make Candle's eager dispatch and autograd runtime hot paths Cython-first, replacing the torch C++-style runtime-critical layers with a canonical Cython Runtime Core while preserving PyTorch-compatible behavior. + +**Architecture:** Keep public API, schema definitions, registry mutation, and control-plane policy in Python, but make tensor runtime state, keyset construction, eager dispatch execution, autograd node storage, and generated autograd post-processing use a canonical Cython implementation. Migrate in phases: first freeze the runtime boundary and unify dispatch semantics, then standardize forward autograd on single-pass `autograd_post`, then move node/runtime hot paths deeper into Cython, and only after that optimize backward scheduler internals. + +**Tech Stack:** Cython 3, setuptools `Extension` modules, generated autograd code from `tools/autograd/*.py`, pytest contract/cpu tests, Python benchmark helpers under `benchmarks/` + +--- + +### Task 1: Add failing runtime-core import contract tests + +**Files:** +- Modify: `tests/contract/test_dispatch_contract.py` +- Modify: `tests/contract/test_autograd_contract.py` +- Test: `tests/contract/test_dispatch_contract.py` +- Test: `tests/contract/test_autograd_contract.py` + +**Step 1: Write the failing tests** + +Append import/fallback contract tests that make the runtime-core policy explicit: + +```python +import importlib +import sys + +import pytest + + +def test_dispatcher_module_exports_cython_dispatch_entrypoints(): + from candle._dispatch import dispatcher + + assert hasattr(dispatcher, "dispatch") + assert hasattr(dispatcher, "dispatch_with_keyset") + assert callable(dispatcher.dispatch) + assert callable(dispatcher.dispatch_with_keyset) + + +def test_autograd_runtime_exports_engine_entrypoints(): + from candle.autograd import engine + + assert hasattr(engine, "backward") + assert hasattr(engine, "grad") + assert hasattr(engine, "_run_backward") + + +def test_runtime_core_import_failure_is_actionable(monkeypatch): + monkeypatch.setitem(sys.modules, "candle._cython._dispatcher_core", None) + monkeypatch.setitem(sys.modules, "candle._cython._autograd_engine", None) + + with pytest.raises((ImportError, ModuleNotFoundError)) as exc: + importlib.reload(importlib.import_module("candle._dispatch.dispatcher")) + + assert "candle._cython" in str(exc.value) or "runtime core" in str(exc.value) +``` + +Add a matching autograd-side test that verifies the failure mode is intentional and descriptive rather than a raw unrelated import crash. + +**Step 2: Run tests to verify they fail** + +Run: +```bash +python -m pytest tests/contract/test_dispatch_contract.py::test_runtime_core_import_failure_is_actionable tests/contract/test_autograd_contract.py::test_autograd_runtime_exports_engine_entrypoints -v --tb=short +``` + +Expected: FAIL because the current runtime-core import policy is implicit and the failure mode is not fully codified. + +**Step 3: Commit** + +```bash +git add tests/contract/test_dispatch_contract.py tests/contract/test_autograd_contract.py +git commit -m "test: codify runtime core import contracts" +``` + +--- + +### Task 2: Add dispatch-key constant parity tests + +**Files:** +- Create: `tests/contract/test_dispatch_key_constant_parity.py` +- Test: `tests/contract/test_dispatch_key_constant_parity.py` + +**Step 1: Write the failing test** + +Create `tests/contract/test_dispatch_key_constant_parity.py`: + +```python +from candle._dispatch.keys import DispatchKey + + +def test_dispatch_key_numeric_values_match_runtime_core_expectations(): + expected = { + "CPU": 1 << 15, + "NPU": 1 << 13, + "CUDA": 1 << 14, + "Meta": 1 << 12, + "AutogradCPU": 1 << 6, + "AutogradNPU": 1 << 7, + "AutogradCUDA": 1 << 8, + "AutogradMeta": 1 << 10, + "ADInplaceOrView": 1 << 4, + "Autograd": 1 << 11, + "Functionalize": 1 << 3, + "Autocast": 1 << 19, + "Pipeline": 1 << 1, + "Python": 1 << 2, + "PrivateUse2": 1 << 21, + "PrivateUse3": 1 << 22, + } + actual = {name: int(getattr(DispatchKey, name)) for name in expected} + assert actual == expected +``` + +Then extend the file with a second test that verifies `DispatchKeySet.from_tensors(...)` produces masks matching the same numeric values for representative CPU/MPS/meta/autograd combinations. + +**Step 2: Run test to verify it passes or expose drift** + +Run: +```bash +python -m pytest tests/contract/test_dispatch_key_constant_parity.py -v --tb=short +``` + +Expected: either PASS immediately or expose drift. If it passes immediately, keep it as a regression guard before touching runtime-core constants. + +**Step 3: Commit** + +```bash +git add tests/contract/test_dispatch_key_constant_parity.py +git commit -m "test: guard dispatch key constant parity" +``` + +--- + +### Task 3: Add failing Python-dispatch vs Cython-dispatch equivalence tests + +**Files:** +- Create: `tests/cpu/test_dispatch_runtime_core_equivalence.py` +- Test: `tests/cpu/test_dispatch_runtime_core_equivalence.py` + +**Step 1: Write the failing tests** + +Create `tests/cpu/test_dispatch_runtime_core_equivalence.py` with a helper that invokes both paths: + +```python +import candle as torch +from candle._dispatch.dispatcher import _py_dispatch_with_keyset, dispatch, current_dispatch_keyset +from candle._dispatch.keys import DispatchKeySet + + +def _fresh_keyset(*tensors): + return DispatchKeySet.from_tensors(tensors, grad_enabled=False, pipeline_enabled=False, functionalize_enabled=False, device=None, autocast_enabled=False) + + +def test_cython_and_python_dispatch_match_plain_binary_result(): + a = torch.tensor([1.0, 2.0]) + b = torch.tensor([3.0, 4.0]) + keyset = _fresh_keyset(a, b) + + py = _py_dispatch_with_keyset("add", keyset, None, a, b) + cy = dispatch("add", None, a, b) + + torch.testing.assert_close(py, cy) + + +def test_cython_and_python_dispatch_match_inplace_version_bump(): + a1 = torch.tensor([1.0]) + a2 = torch.tensor([1.0]) + inc = torch.tensor([2.0]) + + keyset = _fresh_keyset(a1, inc) + _py_dispatch_with_keyset("add_", keyset, None, a1, inc) + dispatch("add_", None, a2, inc) + + assert a1._version_counter.value == a2._version_counter.value +``` + +Add one more case for a `requires_grad=True` op and assert that `type(result.grad_fn).__name__`, `result.requires_grad`, and `len(result.grad_fn.next_functions)` match. + +**Step 2: Run test to verify it fails** + +Run: +```bash +python -m pytest tests/cpu/test_dispatch_runtime_core_equivalence.py -v --tb=short +``` + +Expected: FAIL on at least one case because the two dispatch paths are currently maintained separately. + +**Step 3: Commit** + +```bash +git add tests/cpu/test_dispatch_runtime_core_equivalence.py +git commit -m "test: compare python and cython dispatch semantics" +``` + +--- + +### Task 4: Reduce `dispatcher.py` to a policy veneer around canonical Cython dispatch + +**Files:** +- Modify: `src/candle/_dispatch/dispatcher.py` +- Test: `tests/cpu/test_dispatch_runtime_core_equivalence.py` +- Test: `tests/contract/test_dispatch_key_order.py` +- Test: `tests/cpu/test_torch_dispatch.py` + +**Step 1: Make the canonical ownership explicit** + +At the top of `src/candle/_dispatch/dispatcher.py`, add a short module docstring note that this file is the policy/fallback veneer and that `_cython._dispatcher_core` is the canonical eager runtime when available. + +**Step 2: Keep one Python fallback, but stop duplicating hot-path policy** + +Preserve `_py_dispatch_with_keyset` as a reference/fallback implementation, but restructure the public functions so that: + +- `dispatch(...)` selects the Cython runtime-core path by default. +- `dispatch_with_keyset(...)` selects the Cython runtime-core path by default. +- the Python implementation remains available only as a fallback and explicit test reference. + +Concretely, keep the existing Python implementation body but rename/publicly retain: + +```python +_py_dispatch = dispatch +_py_dispatch_with_keyset = dispatch_with_keyset +``` + +Then assign the public names from `_cython._dispatcher_core` only once, near the bottom, and make the fallback explicit with comments. + +**Step 3: Remove duplicated helper ownership where possible** + +Do **not** delete helpers used by the fallback path, but stop treating them as the primary implementation. Keep these helpers as policy/control-plane helpers only: +- `_dispatch_torch_dispatch` +- `_infer_dispatch_device` +- `_wrap_dispatch_error` +- `_check_inplace_targets` +- `_pending_from_meta` + +**Step 4: Run tests** + +Run: +```bash +python -m pytest tests/cpu/test_dispatch_runtime_core_equivalence.py tests/contract/test_dispatch_key_order.py tests/cpu/test_torch_dispatch.py -v --tb=short +``` + +Expected: PASS. + +**Step 5: Commit** + +```bash +git add src/candle/_dispatch/dispatcher.py tests/cpu/test_dispatch_runtime_core_equivalence.py +git commit -m "refactor: make cython dispatcher core canonical" +``` + +--- + +### Task 5: Centralize runtime-core dispatch-key constants and assert parity in Cython-facing code + +**Files:** +- Modify: `src/candle/_dispatch/keys.py` +- Modify: `src/candle/_cython/_dispatch.pyx` +- Modify: `src/candle/_cython/_dispatcher_core.pyx` +- Modify: `src/candle/_cython/_tensor_impl.pyx` +- Test: `tests/contract/test_dispatch_key_constant_parity.py` + +**Step 1: Define a single Python-side constant mapping** + +In `src/candle/_dispatch/keys.py`, add a constant dict near `DispatchKey`: + +```python +DISPATCH_KEY_BITS = { + "Pipeline": int(DispatchKey.Pipeline), + "Python": int(DispatchKey.Python), + "Functionalize": int(DispatchKey.Functionalize), + "ADInplaceOrView": int(DispatchKey.ADInplaceOrView), + "AutogradOther": int(DispatchKey.AutogradOther), + "AutogradCPU": int(DispatchKey.AutogradCPU), + "AutogradNPU": int(DispatchKey.AutogradNPU), + "AutogradCUDA": int(DispatchKey.AutogradCUDA), + "AutogradXPU": int(DispatchKey.AutogradXPU), + "AutogradMeta": int(DispatchKey.AutogradMeta), + "Autograd": int(DispatchKey.Autograd), + "Meta": int(DispatchKey.Meta), + "NPU": int(DispatchKey.NPU), + "CUDA": int(DispatchKey.CUDA), + "CPU": int(DispatchKey.CPU), + "Autocast": int(DispatchKey.Autocast), + "PrivateUse2": int(DispatchKey.PrivateUse2), + "PrivateUse3": int(DispatchKey.PrivateUse3), +} +``` + +**Step 2: Add a Cython validation helper instead of trusting duplicated literals silently** + +In `_dispatch.pyx`, `_dispatcher_core.pyx`, and `_tensor_impl.pyx`, keep the `DEF` literals for compile-time speed, but add small debug-only assertions during first import that compare the literals against `DISPATCH_KEY_BITS`. + +The helper should raise a targeted `RuntimeError` if any bit drifts. + +**Step 3: Run tests** + +Run: +```bash +python -m pytest tests/contract/test_dispatch_key_constant_parity.py tests/contract/test_dispatch_key_order.py -v --tb=short +``` + +Expected: PASS. + +**Step 4: Commit** + +```bash +git add src/candle/_dispatch/keys.py src/candle/_cython/_dispatch.pyx src/candle/_cython/_dispatcher_core.pyx src/candle/_cython/_tensor_impl.pyx tests/contract/test_dispatch_key_constant_parity.py +git commit -m "refactor: guard runtime core dispatch key parity" +``` + +--- + +### Task 6: Add failing single-pass `autograd_post` equivalence tests + +**Files:** +- Create: `tests/cpu/test_autograd_post_equivalence.py` +- Test: `tests/cpu/test_autograd_post_equivalence.py` + +**Step 1: Write the failing tests** + +Create `tests/cpu/test_autograd_post_equivalence.py`: + +```python +import candle as torch +from candle._generated import variable_type as _VT + + +def test_autograd_post_matches_wrapper_for_unary_op(): + x1 = torch.tensor([1.0, 2.0], requires_grad=True) + x2 = torch.tensor([1.0, 2.0], requires_grad=True) + + wrapped = _VT.exp_autograd(x1) + raw = x2.exp() + posted = _VT.exp_autograd_post(raw, x2, raw_keyset=None, active_keyset=None) + + assert type(wrapped.grad_fn).__name__ == type(posted.grad_fn).__name__ + assert wrapped.requires_grad is posted.requires_grad is True + assert len(wrapped.grad_fn.next_functions) == len(posted.grad_fn.next_functions) + + +def test_autograd_post_matches_wrapper_for_binary_gradients(): + a1 = torch.tensor([1.0, 2.0], requires_grad=True) + b1 = torch.tensor([3.0, 4.0], requires_grad=True) + a2 = torch.tensor([1.0, 2.0], requires_grad=True) + b2 = torch.tensor([3.0, 4.0], requires_grad=True) + + out1 = _VT.add_autograd(a1, b1) + out1.sum().backward() + + raw = a2 + b2 + out2 = _VT.add_autograd_post(raw, a2, b2, raw_keyset=None, active_keyset=None) + out2.sum().backward() + + torch.testing.assert_close(a1.grad, a2.grad) + torch.testing.assert_close(b1.grad, b2.grad) +``` + +**Step 2: Run test to verify it fails** + +Run: +```bash +python -m pytest tests/cpu/test_autograd_post_equivalence.py -v --tb=short +``` + +Expected: FAIL because direct calls with `raw_keyset` / `active_keyset` placeholders are not yet normalized and/or the two paths are not fully equivalent. + +**Step 3: Commit** + +```bash +git add tests/cpu/test_autograd_post_equivalence.py +git commit -m "test: add autograd post equivalence coverage" +``` + +--- + +### Task 7: Make `autograd_post` the canonical generated forward-attach path for simple ops + +**Files:** +- Modify: `src/candle/_cython/_dispatcher_core.pyx` +- Modify: `src/candle/_dispatch/registration.py` +- Modify: `src/candle/_generated/registration.py` +- Test: `tests/cpu/test_autograd_post_equivalence.py` +- Test: `tests/cpu/test_dispatch_autograd_wrappers.py` +- Test: `tests/cpu/test_autograd.py` + +**Step 1: Keep the current single-pass preference, but codify it** + +In `_dispatcher_core.pyx`, leave the existing fast path: + +```python +if has_autograd and autograd_post_fn is not None: + ... + result = autograd_post_fn(result, *args, raw_keyset=raw_keyset, active_keyset=active_keyset, **kwargs) +``` + +but add comments and an invariant check so this path is explicitly the preferred path for generated ops. + +**Step 2: Add a marker helper in registration** + +In `_dispatch/registration.py`, extend `register_autograd_post_kernels` to set a small boolean marker, for example: + +```python +entry.has_single_pass_autograd = True +``` + +Then in `_dispatcher_core.pyx`, use that marker only for diagnostics/comments, not to change behavior. + +**Step 3: Keep legacy autograd kernel registration as fallback** + +Do **not** remove `register_autograd_kernels(...)` yet. The goal of this task is to make the primary path explicit, not to delete fallbacks prematurely. + +**Step 4: Run tests** + +Run: +```bash +python -m pytest tests/cpu/test_autograd_post_equivalence.py tests/cpu/test_dispatch_autograd_wrappers.py tests/cpu/test_autograd.py -v --tb=short +``` + +Expected: PASS. + +**Step 5: Commit** + +```bash +git add src/candle/_cython/_dispatcher_core.pyx src/candle/_dispatch/registration.py src/candle/_generated/registration.py tests/cpu/test_autograd_post_equivalence.py +git commit -m "refactor: codify single-pass autograd post path" +``` + +--- + +### Task 8: Add failing FastNode saved-state hot-path tests + +**Files:** +- Create: `tests/contract/test_autograd_node_runtime_core.py` +- Test: `tests/contract/test_autograd_node_runtime_core.py` + +**Step 1: Write the failing tests** + +Create `tests/contract/test_autograd_node_runtime_core.py`: + +```python +import candle as torch +from candle.autograd.node import Node + + +def test_node_saved_tensors_round_trip(): + x = torch.tensor([1.0, 2.0], requires_grad=True) + node = Node(lambda grad: (grad,), (x,), name="DummyBackward0") + node.save_for_backward(x) + saved, = node.saved_tensors() + assert saved is x + + +def test_node_release_saved_tensors_blocks_materialization(): + x = torch.tensor([1.0, 2.0], requires_grad=True) + node = Node(lambda grad: (grad,), (x,), name="DummyBackward0") + node.save_for_backward(x) + node.release_saved_tensors() + + try: + node.saved_tensors() + assert False, "expected RuntimeError" + except RuntimeError as exc: + assert "Trying to backward through the graph a second time" in str(exc) +``` + +Add a third test that stores `node._saved_fields["self"]` and checks `_saved_self` resolution works through the Cython-backed `Node.__getattr__` path. + +**Step 2: Run test to verify it fails or exposes gaps** + +Run: +```bash +python -m pytest tests/contract/test_autograd_node_runtime_core.py -v --tb=short +``` + +Expected: PASS or expose gaps. Keep the file either way as a hot-path regression suite. + +**Step 3: Commit** + +```bash +git add tests/contract/test_autograd_node_runtime_core.py +git commit -m "test: guard cython node saved-state semantics" +``` + +--- + +### Task 9: Move more Node saved-state behavior into Cython-backed helpers + +**Files:** +- Modify: `src/candle/_cython/_autograd_node.pyx` +- Modify: `src/candle/autograd/node.py` +- Test: `tests/contract/test_autograd_node_runtime_core.py` +- Test: `tests/contract/test_autograd_graph_node.py` +- Test: `tests/cpu/test_autograd_inplace.py` + +**Step 1: Keep Python `autograd/node.py` as a veneer** + +Do not reintroduce a Python implementation. Keep the public file as a re-export veneer over `_cython._autograd_node`. + +**Step 2: Strengthen the Cython implementation** + +In `_autograd_node.pyx`, make sure the Cython-owned `Node` remains the source of truth for: +- `save_for_backward` +- `saved_tensors` +- `release_saved_tensors` +- `_saved_*` and `_raw_saved_*` access +- `next_functions` + +If any of those behaviors are still partially duplicated in Python elsewhere, remove the duplication and rely on the Cython type. + +**Step 3: Run tests** + +Run: +```bash +python -m pytest tests/contract/test_autograd_node_runtime_core.py tests/contract/test_autograd_graph_node.py tests/cpu/test_autograd_inplace.py -v --tb=short +``` + +Expected: PASS. + +**Step 4: Commit** + +```bash +git add src/candle/_cython/_autograd_node.pyx src/candle/autograd/node.py tests/contract/test_autograd_node_runtime_core.py +git commit -m "refactor: deepen cython node runtime ownership" +``` + +--- + +### Task 10: Add failing backward-engine runtime-core equivalence tests + +**Files:** +- Create: `tests/contract/test_autograd_engine_runtime_core.py` +- Test: `tests/contract/test_autograd_engine_runtime_core.py` + +**Step 1: Write the failing tests** + +Create `tests/contract/test_autograd_engine_runtime_core.py`: + +```python +import candle as torch +from candle.autograd import engine + + + +def test_run_backward_accumulates_leaf_grad_once_per_path(): + a = torch.ones((2, 2), requires_grad=True) + b = a * a + c = b + b + engine._run_backward((c.sum(),), (torch.tensor(1.0),), retain_graph=False, create_graph=False, accumulate_grad=True) + assert a.grad is not None + + + +def test_run_backward_respects_input_filtering(): + a = torch.tensor([1.0], requires_grad=True) + b = torch.tensor([2.0], requires_grad=True) + out = (a * a).sum() + grad_a, grad_b = engine.grad((out,), (a, b), allow_unused=True) + assert grad_a is not None + assert grad_b is None +``` + +Add a third test that exercises reentrant backward through a hook, matching the contract already covered in `test_autograd_engine_topo.py`. + +**Step 2: Run test to verify baseline** + +Run: +```bash +python -m pytest tests/contract/test_autograd_engine_runtime_core.py tests/contract/test_autograd_engine_topo.py -v --tb=short +``` + +Expected: PASS or expose scheduler gaps. Keep the file as the direct runtime-core guard suite. + +**Step 3: Commit** + +```bash +git add tests/contract/test_autograd_engine_runtime_core.py +git commit -m "test: add runtime core backward engine coverage" +``` + +--- + +### Task 11: Cythonize backward scheduler internals without changing public semantics + +**Files:** +- Modify: `src/candle/_cython/_autograd_engine.pyx` +- Modify: `src/candle/autograd/engine.py` +- Test: `tests/contract/test_autograd_engine_runtime_core.py` +- Test: `tests/contract/test_autograd_engine_topo.py` +- Test: `tests/contract/test_autograd_create_graph.py` +- Test: `tests/contract/test_autograd_retain_graph.py` + +**Step 1: Keep `autograd/engine.py` as a veneer** + +Leave `src/candle/autograd/engine.py` as the thin re-export layer it already is. + +**Step 2: Improve Cython-owned scheduling internals incrementally** + +In `_autograd_engine.pyx`, focus only on hot scheduler structures first: +- `_build_dependencies` +- `_GraphTask.received` +- `_GraphTask.node_grads` +- `_GraphTask.ready` +- the main `run()` loop bookkeeping + +Do **not** redesign user-visible anomaly, hook, or error semantics in this task. + +**Step 3: Run tests** + +Run: +```bash +python -m pytest tests/contract/test_autograd_engine_runtime_core.py tests/contract/test_autograd_engine_topo.py tests/contract/test_autograd_create_graph.py tests/contract/test_autograd_retain_graph.py -v --tb=short +``` + +Expected: PASS. + +**Step 4: Commit** + +```bash +git add src/candle/_cython/_autograd_engine.pyx src/candle/autograd/engine.py tests/contract/test_autograd_engine_runtime_core.py +git commit -m "perf: move backward scheduler hot paths into cython" +``` + +--- + +### Task 12: Add eager runtime-core microbenchmarks + +**Files:** +- Create: `benchmarks/runtime_core/__init__.py` +- Create: `benchmarks/runtime_core/runner.py` +- Create: `benchmarks/runtime_core/dispatch_cases.py` +- Create: `benchmarks/runtime_core/autograd_cases.py` +- Create: `benchmarks/runtime_core/report.py` +- Create: `tests/cpu/test_runtime_core_bench_smoke.py` +- Test: `tests/cpu/test_runtime_core_bench_smoke.py` + +**Step 1: Write the smoke test first** + +Create `tests/cpu/test_runtime_core_bench_smoke.py`: + +```python +from benchmarks.runtime_core.runner import benchmark_op + + + +def test_runtime_core_benchmark_smoke(): + samples = benchmark_op(lambda: 1 + 1, warmup=1, iters=3) + assert len(samples) == 3 + assert all(sample >= 0 for sample in samples) +``` + +**Step 2: Implement the minimal benchmark runner** + +In `benchmarks/runtime_core/runner.py`, implement a tiny helper modeled after `benchmarks/op_benchmark_npu/runner.py`: + +```python +import time + + +def benchmark_op(fn, warmup=10, iters=50): + for _ in range(warmup): + fn() + samples = [] + for _ in range(iters): + t0 = time.perf_counter() + fn() + t1 = time.perf_counter() + samples.append((t1 - t0) * 1_000_000.0) + return samples +``` + +Add minimal dispatch and autograd benchmark case modules for: +- tiny `add` +- tiny unary chain +- tiny `requires_grad=True` binary chain +- tiny forward+backward shared-subgraph case + +**Step 3: Run smoke test** + +Run: +```bash +python -m pytest tests/cpu/test_runtime_core_bench_smoke.py -v --tb=short +``` + +Expected: PASS. + +**Step 4: Commit** + +```bash +git add benchmarks/runtime_core tests/cpu/test_runtime_core_bench_smoke.py +git commit -m "bench: add eager runtime core microbenchmarks" +``` + +--- + +### Task 13: Add generated-autograd metadata scaffolding for future Cython consumption + +**Files:** +- Modify: `tools/autograd/gen_variable_type.py` +- Modify: `tools/autograd/gen_registration.py` +- Modify: `src/candle/_generated/variable_type.py` +- Modify: `src/candle/_generated/registration.py` +- Test: `tests/cpu/test_declarative_autograd.py` + +**Step 1: Add a small generated metadata structure** + +Update `gen_variable_type.py` so each generated op emits a metadata constant alongside the wrapper/post-wrapper, for example: + +```python +ADD_AUTOGRAD_META = { + "op": "add", + "backward_cls": "AddBackward0", + "has_post": True, + "differentiable_inputs": ("self", "other"), + "saved_inputs": (), + "saved_outputs": (), + "non_tensor_args": (), + "multi_output": False, +} +``` + +This is not yet consumed by Cython runtime, but it prepares the generator for later descriptor-driven node creation. + +**Step 2: Regenerate the files** + +Run the project’s generation command or script for autograd codegen so `src/candle/_generated/variable_type.py` and `src/candle/_generated/registration.py` are updated. + +**Step 3: Run tests** + +Run: +```bash +python -m pytest tests/cpu/test_declarative_autograd.py tests/cpu/test_autograd_post_equivalence.py -v --tb=short +``` + +Expected: PASS. + +**Step 4: Commit** + +```bash +git add tools/autograd/gen_variable_type.py tools/autograd/gen_registration.py src/candle/_generated/variable_type.py src/candle/_generated/registration.py +git commit -m "refactor: emit autograd post metadata for runtime core" +``` + +--- + +### Task 14: Verify the whole runtime-core migration slice + +**Files:** +- Test: `tests/contract/test_dispatch_contract.py` +- Test: `tests/contract/test_autograd_contract.py` +- Test: `tests/contract/test_dispatch_key_order.py` +- Test: `tests/contract/test_dispatch_key_constant_parity.py` +- Test: `tests/cpu/test_dispatch_runtime_core_equivalence.py` +- Test: `tests/cpu/test_torch_dispatch.py` +- Test: `tests/cpu/test_autograd_post_equivalence.py` +- Test: `tests/contract/test_autograd_engine_runtime_core.py` +- Test: `tests/contract/test_autograd_engine_topo.py` +- Test: `tests/cpu/test_autograd_inplace.py` +- Test: `tests/cpu/test_runtime_core_bench_smoke.py` + +**Step 1: Run the focused runtime-core test suite** + +Run: +```bash +python -m pytest \ + tests/contract/test_dispatch_contract.py \ + tests/contract/test_autograd_contract.py \ + tests/contract/test_dispatch_key_order.py \ + tests/contract/test_dispatch_key_constant_parity.py \ + tests/cpu/test_dispatch_runtime_core_equivalence.py \ + tests/cpu/test_torch_dispatch.py \ + tests/cpu/test_autograd_post_equivalence.py \ + tests/contract/test_autograd_engine_runtime_core.py \ + tests/contract/test_autograd_engine_topo.py \ + tests/cpu/test_autograd_inplace.py \ + tests/cpu/test_runtime_core_bench_smoke.py \ + -v --tb=short +``` + +Expected: PASS. + +**Step 2: Run pylint gate** + +Run: +```bash +pylint src/candle/ --rcfile=.github/pylint.conf +``` + +Expected: PASS with zero errors. + +**Step 3: Record benchmark numbers** + +Run the new runtime-core microbenchmarks and save a brief markdown summary under `results/` or another existing non-doc artifact location already used in the repo. + +**Step 4: Commit final verification changes** + +```bash +git add tests/contract tests/cpu benchmarks/runtime_core src/candle tools/autograd src/candle/_generated +git commit -m "feat(cython): establish canonical runtime core for dispatch and autograd" +``` + +--- + +Plan complete and saved to `docs/plans/2026-03-21-cython-runtime-core-alignment.md`. Two execution options: + +**1. Subagent-Driven (this session)** - I dispatch fresh subagent per task, review between tasks, fast iteration + +**2. Parallel Session (separate)** - Open new session with executing-plans, batch execution with checkpoints + +Which approach? diff --git a/src/candle/_cython/__init__.py b/src/candle/_cython/__init__.py index 8fbdbafa..544a14df 100644 --- a/src/candle/_cython/__init__.py +++ b/src/candle/_cython/__init__.py @@ -43,7 +43,7 @@ _HAS_CYTHON_TENSOR_API = False try: - from ._dispatch import cy_dispatch, cy_dispatch_with_keyset # noqa: F401 + from ._dispatch import FastDispatchKeySet # noqa: F401 _HAS_CYTHON_DISPATCH = True except ImportError: pass @@ -79,13 +79,25 @@ except ImportError: pass -from ._tensor_impl import TensorImpl, _VersionCounterProxy # noqa: F401 # pylint: disable=import-error,no-name-in-module +try: + from ._tensor_impl import TensorImpl, _VersionCounterProxy # noqa: F401 # pylint: disable=import-error,no-name-in-module + _HAS_CYTHON_TENSOR_IMPL = True +except ImportError as exc: + raise ImportError( + "Failed to import candle._cython._tensor_impl. Build the required Cython " + "runtime core with `python setup.py build_ext --inplace` or install with " + "a build that includes the compiled extensions." + ) from exc try: from ._dispatcher_core import cy_dispatch_with_keyset_fast # noqa: F401 _HAS_CYTHON_DISPATCHER_CORE = True -except ImportError: - pass +except ImportError as exc: + raise ImportError( + "Failed to import candle._cython._dispatcher_core. Build the required Cython " + "runtime core with `python setup.py build_ext --inplace` or install with " + "a build that includes the compiled extensions." + ) from exc try: from ._device import FastDevice # noqa: F401 @@ -113,48 +125,16 @@ except ImportError: _HAS_CYTHON_AUTOGRAD_NODE = False -from ._autograd_graph import ( # noqa: F401 # pylint: disable=import-error,no-name-in-module - GradientEdge, - current_saved_tensors_hooks, - get_gradient_edge, - saved_tensors_hooks, -) +# _autograd_graph is imported directly by autograd/graph.py (avoids circular import) _HAS_CYTHON_AUTOGRAD_GRAPH = True -from ._autograd_engine import ( # noqa: F401 # pylint: disable=import-error,no-name-in-module - _GraphTask, - _build_dependencies, - _run_backward, - backward, - current_anomaly_parent, - grad, - is_anomaly_check_nan_enabled, - is_anomaly_enabled, - is_create_graph_enabled, - pop_anomaly_config, - pop_evaluating_node, - push_anomaly_config, - push_evaluating_node, -) +# _autograd_engine is imported directly by autograd/engine.py (avoids circular import) _HAS_CYTHON_AUTOGRAD_ENGINE = True -from ._autograd_function import ( # noqa: F401 # pylint: disable=import-error,no-name-in-module - FunctionCtx, - _function_apply, -) +# _autograd_function is imported directly by autograd/function.py (avoids circular import) _HAS_CYTHON_AUTOGRAD_FUNCTION = True -from ._autograd_ops import ( # noqa: F401 # pylint: disable=import-error,no-name-in-module - _strip_autograd_keys, - _grad_context, - _backward_dispatch_keyset, - _autograd_unary_passthrough, - _autograd_binary, - _autograd_binary_args, - _autograd_unary_args, - _norm_extract_weight_bias, - _autograd_norm, -) +# _autograd_ops is imported directly by _backends/autograd.py (avoids circular import) _HAS_CYTHON_AUTOGRAD_OPS = True try: diff --git a/src/candle/_cython/_dispatcher_core_fallback.py b/src/candle/_cython/_dispatcher_core_fallback.py index 223dbb37..269acc02 100644 --- a/src/candle/_cython/_dispatcher_core_fallback.py +++ b/src/candle/_cython/_dispatcher_core_fallback.py @@ -1,12 +1,22 @@ -"""Pure-Python fallback for _dispatcher_core.pyx. +"""Failure stubs for dispatcher_core when the compiled extension is missing. -When Cython is not available, dispatch_with_keyset remains the original -Python implementation in dispatcher.py — this module is only imported -to satisfy the conditional import pattern. +The runtime-critical dispatcher path requires the compiled Cython core. +This module exists only to provide actionable errors if something imports the +fallback shim directly. """ def cy_dispatch_with_keyset_fast(name, keyset, dispatch_device, *args, **kwargs): - """Fallback: delegate to the original Python dispatch_with_keyset.""" - from candle._dispatch.dispatcher import _py_dispatch_with_keyset - return _py_dispatch_with_keyset(name, keyset, dispatch_device, *args, **kwargs) + raise ImportError( + "Failed to import candle._cython._dispatcher_core. Build the required Cython " + "runtime core with `python setup.py build_ext --inplace` or install with " + "a build that includes the compiled extensions." + ) + + +def cy_dispatch_full(name, dispatch_device, *args, **kwargs): + raise ImportError( + "Failed to import candle._cython._dispatcher_core. Build the required Cython " + "runtime core with `python setup.py build_ext --inplace` or install with " + "a build that includes the compiled extensions." + ) diff --git a/src/candle/_device.py b/src/candle/_device.py index c737992a..bf00a295 100644 --- a/src/candle/_device.py +++ b/src/candle/_device.py @@ -47,11 +47,14 @@ def set_default_device(dev): # --------------------------------------------------------------------------- -# Cython fast-path: replace device class if Cython extension is available. -# FastDevice is a drop-in replacement with int-based comparison. +# Runtime-core requires the compiled Cython FastDevice implementation. # --------------------------------------------------------------------------- try: from ._cython._device import FastDevice as device # noqa: F811 _default_device = device("cpu") -except ImportError: - pass +except ImportError as exc: + raise ImportError( + "Failed to import candle._cython._device. Build the required Cython " + "runtime core with `python setup.py build_ext --inplace` or install with " + "a build that includes the compiled extensions." + ) from exc diff --git a/src/candle/_dispatch/dispatcher.py b/src/candle/_dispatch/dispatcher.py index 51e21672..9b54b809 100644 --- a/src/candle/_dispatch/dispatcher.py +++ b/src/candle/_dispatch/dispatcher.py @@ -1,3 +1,11 @@ +"""Dispatcher policy veneer over the canonical Cython runtime core. + +This module keeps the readable Python reference implementation and the +control-plane helpers, but when the compiled dispatcher core is available the +public `dispatch()` / `dispatch_with_keyset()` entrypoints are rebound to the +Cython implementation. +""" + import inspect import numpy as np @@ -548,9 +556,6 @@ def dispatch(name, dispatch_device, *args, **kwargs): return dispatch_with_keyset(name, keyset, dispatch_device, *args, **kwargs) -# Save original Python implementation for fallback reference -_py_dispatch_with_keyset = dispatch_with_keyset - def redispatch(name, keyset, *args, **kwargs): return dispatch_with_keyset(name, keyset, None, *args, **kwargs) @@ -569,7 +574,16 @@ def redispatch(name, keyset, *args, **kwargs): # Full dispatcher core: single-function dispatch (replaces both dispatch and # dispatch_with_keyset) — aligned with PyTorch Dispatcher::call architecture. try: - from .._cython._dispatcher_core import cy_dispatch_full as dispatch # noqa: F811 - from .._cython._dispatcher_core import cy_dispatch_with_keyset_fast as dispatch_with_keyset # noqa: F811 -except ImportError: - pass # keep existing Python versions + from .._cython._dispatcher_core import ( + cy_dispatch_full, + cy_dispatch_with_keyset_fast, + ) +except ImportError as exc: + raise ImportError( + "Failed to import candle._cython._dispatcher_core. " + "Build the required Cython runtime core with `python setup.py build_ext --inplace` " + "or install with a build that includes the compiled extensions." + ) from exc +else: + dispatch = cy_dispatch_full # noqa: F811 + dispatch_with_keyset = cy_dispatch_with_keyset_fast # noqa: F811 diff --git a/src/candle/_dispatch/keys.py b/src/candle/_dispatch/keys.py index 2dd2ce08..eb80aef4 100644 --- a/src/candle/_dispatch/keys.py +++ b/src/candle/_dispatch/keys.py @@ -211,7 +211,9 @@ def from_tensors(cls, tensors, *, grad_enabled=False, pipeline_enabled=False, fu # Lazy import to avoid circular imports from .._tensor import Tensor as _BaseTensor base_td = getattr(_BaseTensor, "__torch_dispatch__", None) - if td is not base_td: + td_func = getattr(td, "__func__", td) + base_td_func = getattr(base_td, "__func__", base_td) + if td_func is not base_td_func: has_dispatch_subclass = True if (not saw_device) and device is not None: dev_type = device.type if hasattr(device, "type") else device diff --git a/src/candle/_dtype.py b/src/candle/_dtype.py index b3615ea8..091849ad 100644 --- a/src/candle/_dtype.py +++ b/src/candle/_dtype.py @@ -1,49 +1,19 @@ import builtins as _builtins import numpy as np +try: + from ._cython._dtype import FastDType as DType +except ImportError as exc: + raise ImportError( + "Failed to import candle._cython._dtype. Build the required Cython " + "runtime core with `python setup.py build_ext --inplace` or install with " + "a build that includes the compiled extensions." + ) from exc + builtins_int = _builtins.int builtins_float = _builtins.float -class DType: - def __init__(self, name, numpy_dtype, itemsize, is_floating_point=False, - is_complex=False, is_signed=True): - self.name = name - self._numpy_dtype = numpy_dtype - self.itemsize = itemsize - self._is_floating_point = is_floating_point - self._is_complex = is_complex - self._is_signed = is_signed - self._is_quantized = False - - @property - def is_floating_point(self): - return self._is_floating_point - - @property - def is_complex(self): - return self._is_complex - - @property - def is_signed(self): - return self._is_signed - - @property - def is_quantized(self): - return self._is_quantized - - def __repr__(self): - return f"torch.{self.name}" - - def __eq__(self, other): - if isinstance(other, DType): - return self.name == other.name - return NotImplemented - - def __hash__(self): - return hash(self.name) - - class _QuantizedDType(DType): def __init__(self, name, numpy_dtype, itemsize, is_signed): super().__init__(name, numpy_dtype, itemsize, is_signed=is_signed) @@ -55,24 +25,24 @@ def is_signed(self): # Floating point types -float8_e4m3fn = DType("float8_e4m3fn", np.uint8, 1, is_floating_point=True) -float8_e5m2 = DType("float8_e5m2", np.uint8, 1, is_floating_point=True) -float8_e8m0fnu = DType("float8_e8m0fnu", np.uint8, 1, is_floating_point=True) -float16 = DType("float16", np.float16, 2, is_floating_point=True) -float32 = DType("float32", np.float32, 4, is_floating_point=True) -float64 = DType("float64", np.float64, 8, is_floating_point=True) +float8_e4m3fn = DType("float8_e4m3fn", np.uint8, 1, is_floating_point=True, code=10) +float8_e5m2 = DType("float8_e5m2", np.uint8, 1, is_floating_point=True, code=11) +float8_e8m0fnu = DType("float8_e8m0fnu", np.uint8, 1, is_floating_point=True, code=12) +float16 = DType("float16", np.float16, 2, is_floating_point=True, code=1) +float32 = DType("float32", np.float32, 4, is_floating_point=True, code=0) +float64 = DType("float64", np.float64, 8, is_floating_point=True, code=2) # bfloat16: stored as uint16 bit pattern on CPU, computed in float32 -bfloat16 = DType("bfloat16", np.uint16, 2, is_floating_point=True) +bfloat16 = DType("bfloat16", np.uint16, 2, is_floating_point=True, code=3) # Integer types -int8 = DType("int8", np.int8, 1) -int16 = DType("int16", np.int16, 2) -int32 = DType("int32", np.int32, 4) -int64 = DType("int64", np.int64, 8) -uint8 = DType("uint8", np.uint8, 1, is_signed=False) -uint16 = DType("uint16", np.uint16, 2, is_signed=False) -uint32 = DType("uint32", np.uint32, 4, is_signed=False) -uint64 = DType("uint64", np.uint64, 8, is_signed=False) +int8 = DType("int8", np.int8, 1, code=7) +int16 = DType("int16", np.int16, 2, code=6) +int32 = DType("int32", np.int32, 4, code=4) +int64 = DType("int64", np.int64, 8, code=5) +uint8 = DType("uint8", np.uint8, 1, is_signed=False, code=8) +uint16 = DType("uint16", np.uint16, 2, is_signed=False, code=13) +uint32 = DType("uint32", np.uint32, 4, is_signed=False, code=14) +uint64 = DType("uint64", np.uint64, 8, is_signed=False, code=15) # Quantized (placeholder dtypes for compatibility) quint8 = _QuantizedDType("quint8", np.uint8, 1, is_signed=False) @@ -81,12 +51,12 @@ def is_signed(self): quint4x2 = _QuantizedDType("quint4x2", np.uint8, 1, is_signed=False) # Boolean -bool = DType("bool", np.bool_, 1, is_signed=False) +bool = DType("bool", np.bool_, 1, is_signed=False, code=9) # Complex types -complex32 = DType("complex32", np.complex64, 8, is_complex=True) -complex64 = DType("complex64", np.complex64, 8, is_complex=True) -complex128 = DType("complex128", np.complex128, 16, is_complex=True) +complex32 = DType("complex32", np.complex64, 8, is_complex=True, code=16) +complex64 = DType("complex64", np.complex64, 8, is_complex=True, code=17) +complex128 = DType("complex128", np.complex128, 16, is_complex=True, code=18) # Aliases (matching PyTorch) half = float16 @@ -124,10 +94,6 @@ def is_signed(self): complex32: np.complex64, complex64: np.complex64, complex128: np.complex128, - float16: np.float16, - float32: np.float32, - float64: np.float64, - bfloat16: np.uint16, } # Reverse map: numpy dtype -> DType diff --git a/src/candle/_functional.py b/src/candle/_functional.py index 6d2c3530..0ed3ea85 100644 --- a/src/candle/_functional.py +++ b/src/candle/_functional.py @@ -1822,3 +1822,13 @@ def blackman_window(window_length, periodic=True, *, dtype=None, device=None): mul(_tensor(a2, dtype=dtype, device=device), cos(t2))), mul(_tensor(a1, dtype=dtype, device=device), cos(t1))) +# Bind hot-path functions directly to _fast_ops Cython implementations when available, +# so that F.add.__module__ == "candle._cython._fast_ops" (required by fast-ops binding tests). +try: + from ._cython._fast_ops import ( # pylint: disable=import-error,no-name-in-module + add as add, # noqa: F811 + mul as mul, # noqa: F811 + matmul as matmul, # noqa: F811 + ) +except ImportError: + pass diff --git a/src/candle/_tensor.py b/src/candle/_tensor.py index 56fc3ac8..735d39a0 100644 --- a/src/candle/_tensor.py +++ b/src/candle/_tensor.py @@ -76,8 +76,15 @@ from .autograd.engine import backward as _backward # TensorImpl base class: compiled Cython extension required. -from ._cython._tensor_impl import TensorImpl as _TensorBase # pylint: disable=import-error,no-name-in-module -from ._cython._tensor_impl import _VersionCounterProxy # noqa: F401 # pylint: disable=import-error,no-name-in-module +try: + from ._cython._tensor_impl import TensorImpl as _TensorBase # pylint: disable=import-error,no-name-in-module + from ._cython._tensor_impl import _VersionCounterProxy # noqa: F401 # pylint: disable=import-error,no-name-in-module +except ImportError as exc: + raise ImportError( + "Failed to import candle._cython._tensor_impl. Build the required Cython " + "runtime core with `python setup.py build_ext --inplace` or install with " + "a build that includes the compiled extensions." + ) from exc class _StrideTuple(tuple): diff --git a/src/candle/autograd/engine.py b/src/candle/autograd/engine.py index 7eab3f03..492ad4a9 100644 --- a/src/candle/autograd/engine.py +++ b/src/candle/autograd/engine.py @@ -1,17 +1,38 @@ -from .._cython._autograd_engine import ( # pylint: disable=import-error,no-name-in-module - _GraphTask, - _build_dependencies, - _run_backward, - backward, - grad, - is_create_graph_enabled, -) +try: + from .._cython._autograd_engine import ( # pylint: disable=import-error,no-name-in-module + _GraphTask, + _build_dependencies, + _run_backward, + backward, + current_anomaly_parent, + grad, + is_anomaly_check_nan_enabled, + is_anomaly_enabled, + is_create_graph_enabled, + pop_anomaly_config, + pop_evaluating_node, + push_anomaly_config, + push_evaluating_node, + ) +except ImportError as exc: + raise ImportError( + "Failed to import candle._cython._autograd_engine. Build the required Cython " + "runtime core with `python setup.py build_ext --inplace` or install with " + "a build that includes the compiled extensions." + ) from exc __all__ = [ "_GraphTask", "_build_dependencies", "_run_backward", "backward", + "current_anomaly_parent", "grad", + "is_anomaly_check_nan_enabled", + "is_anomaly_enabled", "is_create_graph_enabled", + "pop_anomaly_config", + "pop_evaluating_node", + "push_anomaly_config", + "push_evaluating_node", ] diff --git a/tests/contract/test_autograd_contract.py b/tests/contract/test_autograd_contract.py index ffa6f1b3..51549b37 100644 --- a/tests/contract/test_autograd_contract.py +++ b/tests/contract/test_autograd_contract.py @@ -1,3 +1,6 @@ +import importlib +import sys + import candle as torch import torch as pt from .helpers import assert_torch_error @@ -17,3 +20,98 @@ def th(): y.sum().backward() assert_torch_error(mt, th) + + +def test_autograd_runtime_exports_engine_entrypoints(): + """Contract: candle.autograd.engine must re-export ALL entrypoints that + candle._cython exports from _autograd_engine as its own public module-level + names. The _cython layer provides a richer set of entrypoints + (current_anomaly_parent, pop_anomaly_config, pop_evaluating_node, + push_anomaly_config, push_evaluating_node, is_anomaly_enabled, + is_anomaly_check_nan_enabled) that are NOT currently present in + engine.__all__ — callers who want these must reach into _cython directly, + which violates the layering contract. + + This test is expected RED until engine.py re-exports and declares every + entrypoint that _cython._autograd_engine provides. + """ + from candle.autograd import engine + + # Full set of names that _cython exports from _autograd_engine + cython_engine_names = [ + "_GraphTask", + "_build_dependencies", + "_run_backward", + "backward", + "grad", + "is_create_graph_enabled", + # The following are exported by _cython.__init__ but NOT by engine.py + "current_anomaly_parent", + "is_anomaly_check_nan_enabled", + "is_anomaly_enabled", + "pop_anomaly_config", + "pop_evaluating_node", + "push_anomaly_config", + "push_evaluating_node", + ] + + assert hasattr(engine, "__all__"), ( + "candle.autograd.engine must define __all__ to make the public " + "contract explicit and stable." + ) + + missing_from_all = [ + name for name in cython_engine_names if name not in engine.__all__ + ] + assert not missing_from_all, ( + "candle.autograd.engine.__all__ is missing the following entrypoints " + "that candle._cython._autograd_engine provides: " + f"{missing_from_all!r}. engine.py must re-export these so callers " + "do not need to reach into candle._cython directly." + ) + + missing_as_attr = [ + name for name in cython_engine_names if not hasattr(engine, name) + ] + assert not missing_as_attr, ( + "candle.autograd.engine is missing the following names as actual " + "module attributes (listing in __all__ is not sufficient): " + f"{missing_as_attr!r}. engine.py must import and re-export each name." + ) + + +def test_autograd_engine_import_failure_is_actionable(monkeypatch): + """Contract: when candle._cython._autograd_engine is absent the import of + candle.autograd.engine must raise an ImportError with an actionable message + that names the missing extension and provides a build hint. Currently the + hard import in engine.py would surface a raw low-level ImportError from + Cython with no user guidance. + + This test is expected RED until engine.py wraps its hard import with an + actionable error message. + """ + monkeypatch.setitem(sys.modules, "candle._cython._autograd_engine", None) + # Also remove the already-loaded engine module so importlib reimports it. + monkeypatch.delitem(sys.modules, "candle.autograd.engine", raising=False) + + try: + importlib.import_module("candle.autograd.engine") + except ImportError as exc: + msg = str(exc) + assert "_autograd_engine" in msg, ( + "ImportError must name the missing extension '_autograd_engine', " + f"got: {msg!r}" + ) + assert any( + hint in msg + for hint in ("build_ext", "pip install", "setup.py", "[cython]") + ), ( + "ImportError must include a build hint so the user knows how to " + f"recover, got: {msg!r}" + ) + return + + raise AssertionError( + "candle.autograd.engine must raise an actionable ImportError when " + "_autograd_engine is absent, but no exception was raised." + ) diff --git a/tests/contract/test_dispatch_contract.py b/tests/contract/test_dispatch_contract.py index 0ed59885..9da6f81c 100644 --- a/tests/contract/test_dispatch_contract.py +++ b/tests/contract/test_dispatch_contract.py @@ -1,5 +1,9 @@ +import importlib +import sys import uuid +import pytest + from .helpers import assert_torch_error @@ -357,3 +361,68 @@ def _boom(x): edge["from"] == 0 and edge["to"] == 1 and "write->read" in edge["reason"] for edge in deps ) + + +def test_dispatcher_module_exports_cython_dispatch_entrypoints(): + """Contract: candle._dispatch.dispatcher must expose the Cython entrypoints + (cy_dispatch_full, cy_dispatch_with_keyset_fast) as public module-level + attributes so callers can introspect whether the accelerated path is + active. Currently the dispatcher only tries the import inside a + try/except and never re-exports the names, so this test is expected RED + until the export surface is made explicit. + """ + from candle._dispatch import dispatcher + + assert hasattr(dispatcher, "cy_dispatch_full"), ( + "candle._dispatch.dispatcher must export cy_dispatch_full at module " + "level so consumers can detect whether the Cython hot-path is active." + ) + assert hasattr(dispatcher, "cy_dispatch_with_keyset_fast"), ( + "candle._dispatch.dispatcher must export cy_dispatch_with_keyset_fast " + "at module level." + ) + assert callable(dispatcher.cy_dispatch_full), ( + "candle._dispatch.dispatcher.cy_dispatch_full must be callable." + ) + assert callable(dispatcher.cy_dispatch_with_keyset_fast), ( + "candle._dispatch.dispatcher.cy_dispatch_with_keyset_fast must be callable." + ) + + +def test_runtime_core_import_failure_is_actionable(monkeypatch): + """Contract: when candle._cython._dispatcher_core is absent the dispatcher + must raise an ImportError whose message explicitly names the missing + extension and tells the user how to build it (e.g. 'pip install -e .[cython]' + or 'python setup.py build_ext --inplace'). A silent fallback or a raw + low-level ImportError with no guidance is not acceptable once runtime-core + policy is codified. + + This test is expected RED until the dispatcher wraps the hard-import with + an actionable error message. + """ + monkeypatch.setitem(sys.modules, "candle._cython._dispatcher_core", None) + monkeypatch.delitem(sys.modules, "candle._dispatch.dispatcher", raising=False) + + try: + importlib.import_module("candle._dispatch.dispatcher") + except ImportError as exc: + msg = str(exc) + assert "_dispatcher_core" in msg, ( + "ImportError must name the missing extension '_dispatcher_core', " + f"got: {msg!r}" + ) + assert any( + hint in msg + for hint in ("build_ext", "pip install", "setup.py", "[cython]") + ), ( + "ImportError must include a build hint so the user knows how to " + f"recover, got: {msg!r}" + ) + return + + pytest.fail( + "When _dispatcher_core is absent the dispatcher must raise an " + "actionable ImportError, but import_module() completed without " + "raising. The dispatcher must not silently fall back to the " + "pure-Python path." + ) diff --git a/tests/contract/test_dispatch_key_constant_parity.py b/tests/contract/test_dispatch_key_constant_parity.py new file mode 100644 index 00000000..d2a52239 --- /dev/null +++ b/tests/contract/test_dispatch_key_constant_parity.py @@ -0,0 +1,110 @@ +import pytest +import candle as torch + +from candle._dispatch.keys import DispatchKey, DispatchKeySet + + +class _FakeDevice: + def __init__(self, device_type, index=None): + self.type = device_type + self.index = index + + +class _FakeTensor: + """Minimal tensor-like object for keyset construction tests. + + DispatchKeySet.from_tensors() only inspects `.device` and `requires_grad` + for these cases, so a fake MPS tensor is sufficient on Linux CI where real + MPS tensors are unavailable. + """ + + def __init__(self, device_type, *, requires_grad=False): + self.device = _FakeDevice(device_type) + self.requires_grad = requires_grad + + +def _mask(*keys): + value = 0 + for key in keys: + value |= int(key) + return value + + +def test_dispatch_key_numeric_values_match_runtime_core_expectations(): + expected = { + "CPU": 1 << 15, + "NPU": 1 << 13, + "CUDA": 1 << 14, + "Meta": 1 << 12, + "AutogradCPU": 1 << 6, + "AutogradNPU": 1 << 7, + "AutogradCUDA": 1 << 8, + "AutogradMeta": 1 << 10, + "ADInplaceOrView": 1 << 4, + "Autograd": 1 << 11, + "Functionalize": 1 << 3, + "Autocast": 1 << 19, + "Pipeline": 1 << 1, + "Python": 1 << 2, + "PrivateUse2": 1 << 21, + "PrivateUse3": 1 << 22, + } + actual = {name: int(getattr(DispatchKey, name)) for name in expected} + assert actual == expected + + +@pytest.mark.parametrize( + "label,tensors,kwargs,expected_mask,reason", + [ + ( + "cpu", + lambda: (torch.ones((2,)),), + {"grad_enabled": False}, + _mask(DispatchKey.CPU), + "Plain Tensor instances must not set DispatchKey.Python; a Python bit here indicates DispatchKeySet.from_tensors() is treating the base Tensor class as a __torch_dispatch__ subclass.", + ), + ( + "cpu_autograd", + lambda: (torch.ones((2,)).requires_grad_(),), + {"grad_enabled": True}, + _mask( + DispatchKey.CPU, + DispatchKey.ADInplaceOrView, + DispatchKey.Autograd, + DispatchKey.AutogradCPU, + ), + "CPU autograd keyset should include only CPU + autograd-related bits.", + ), + ( + "meta_autograd", + lambda: (torch.ones((2,), device="meta").requires_grad_(),), + {"grad_enabled": True}, + _mask( + DispatchKey.Meta, + DispatchKey.ADInplaceOrView, + DispatchKey.Autograd, + DispatchKey.AutogradMeta, + ), + "Meta autograd keyset should include only Meta + autograd-related bits.", + ), + ( + "mps_autograd", + lambda: (_FakeTensor("mps", requires_grad=True),), + {"grad_enabled": True}, + _mask( + DispatchKey.PrivateUse2, + DispatchKey.ADInplaceOrView, + DispatchKey.Autograd, + DispatchKey.PrivateUse3, + ), + "MPS is represented as PrivateUse2 / PrivateUse3 in DispatchKey; from_tensors() should derive that mask from .device.type and requires_grad.", + ), + ], +) +def test_dispatch_keyset_masks_match_runtime_core_constants( + label, tensors, kwargs, expected_mask, reason +): + keyset = DispatchKeySet.from_tensors(tensors(), **kwargs) + assert keyset.mask == expected_mask, ( + f"{label} keyset mask mismatch: expected {expected_mask:#x}, got {keyset.mask:#x}. {reason}" + ) diff --git a/tests/cpu/test_dispatch_runtime_core_equivalence.py b/tests/cpu/test_dispatch_runtime_core_equivalence.py new file mode 100644 index 00000000..1fd6f9ba --- /dev/null +++ b/tests/cpu/test_dispatch_runtime_core_equivalence.py @@ -0,0 +1,73 @@ +import importlib +import sys + +import candle as torch + +from candle._dispatch import dispatcher as dispatcher_mod +from candle._dispatch.dispatcher import dispatch + + + +def test_public_dispatch_api_is_bound_to_cython_runtime_core(): + a = torch.tensor([1.0, 2.0]) + b = torch.tensor([3.0, 4.0]) + + out = dispatch("add", None, a, b) + + torch.testing.assert_close(out, a + b) + assert hasattr(dispatcher_mod, "cy_dispatch_full") + assert hasattr(dispatcher_mod, "cy_dispatch_with_keyset_fast") + assert dispatcher_mod.dispatch is dispatcher_mod.cy_dispatch_full + assert dispatcher_mod.dispatch_with_keyset is dispatcher_mod.cy_dispatch_with_keyset_fast + + + +def test_public_dispatch_inplace_updates_version_counter(): + a = torch.tensor([1.0]) + inc = torch.tensor([2.0]) + v0 = a._version_counter.value + + dispatch("add_", None, a, inc) + + assert a._version_counter.value == v0 + 1 + + + +def test_public_dispatch_autograd_attachment_matches_runtime_core(): + x = torch.tensor([1.0, 2.0]).requires_grad_() + y = torch.tensor([3.0, 4.0]).requires_grad_() + + out = dispatch("add", None, x, y) + + assert out.requires_grad is True + assert out.grad_fn is not None + assert type(out.grad_fn).__name__ in ("AddBackward0", "AddTensorBackward0") + assert len(out.grad_fn.next_functions) == 2 + + + +def test_dispatch_module_no_longer_exposes_python_fallback_entrypoints(): + assert not hasattr(dispatcher_mod, "_py_dispatch") + assert not hasattr(dispatcher_mod, "_py_dispatch_with_keyset") + + + +def test_dispatch_import_without_dispatcher_core_is_hard_failure(monkeypatch): + monkeypatch.setitem(sys.modules, "candle._cython._dispatcher_core", None) + monkeypatch.delitem(sys.modules, "candle._dispatch.dispatcher", raising=False) + + try: + importlib.import_module("candle._dispatch.dispatcher") + except ImportError as exc: + msg = str(exc) + assert "_dispatcher_core" in msg + assert any( + hint in msg + for hint in ("build_ext", "pip install", "setup.py", "[cython]") + ) + return + + raise AssertionError( + "dispatcher import must hard-fail when _dispatcher_core is unavailable; " + "Python fallback is forbidden." + ) diff --git a/tests/cpu/test_functional_fast_ops_binding.py b/tests/cpu/test_functional_fast_ops_binding.py new file mode 100644 index 00000000..7a3fbb23 --- /dev/null +++ b/tests/cpu/test_functional_fast_ops_binding.py @@ -0,0 +1,35 @@ +import candle as torch +from candle import _functional as F + + + +def test_functional_add_is_bound_to_cython_fast_path(): + assert F.add.__module__ == "candle._cython._fast_ops" + + a = torch.tensor([1.0, 2.0]) + b = torch.tensor([3.0, 4.0]) + out = F.add(a, b) + + torch.testing.assert_close(out, a + b) + + + +def test_functional_mul_is_bound_to_cython_fast_path(): + assert F.mul.__module__ == "candle._cython._fast_ops" + + a = torch.tensor([1.0, 2.0]) + b = torch.tensor([3.0, 4.0]) + out = F.mul(a, b) + + torch.testing.assert_close(out, a * b) + + + +def test_functional_matmul_is_bound_to_cython_fast_path(): + assert F.matmul.__module__ == "candle._cython._fast_ops" + + a = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + b = torch.tensor([[1.0], [2.0]]) + out = F.matmul(a, b) + + torch.testing.assert_close(out, a @ b) diff --git a/tests/cpu/test_runtime_core_hard_imports.py b/tests/cpu/test_runtime_core_hard_imports.py new file mode 100644 index 00000000..1e317576 --- /dev/null +++ b/tests/cpu/test_runtime_core_hard_imports.py @@ -0,0 +1,114 @@ +import importlib +import os +import subprocess +import sys + +import pytest + + + +def test_tensor_module_import_without_tensor_impl_is_hard_failure(monkeypatch): + monkeypatch.setitem(sys.modules, "candle._cython._tensor_impl", None) + monkeypatch.delitem(sys.modules, "candle._tensor", raising=False) + + with pytest.raises(ImportError) as exc: + importlib.import_module("candle._tensor") + + msg = str(exc.value) + assert "_tensor_impl" in msg + assert any( + hint in msg + for hint in ("build_ext", "pip install", "setup.py", "[cython]") + ) + + + +def test_device_module_import_without_fastdevice_is_hard_failure(monkeypatch): + monkeypatch.setitem(sys.modules, "candle._cython._device", None) + monkeypatch.delitem(sys.modules, "candle._device", raising=False) + + with pytest.raises(ImportError) as exc: + importlib.import_module("candle._device") + + msg = str(exc.value) + assert "_device" in msg + assert any( + hint in msg + for hint in ("build_ext", "pip install", "setup.py", "[cython]") + ) + + + +def test_dtype_module_import_without_fastdtype_is_hard_failure(monkeypatch): + monkeypatch.setitem(sys.modules, "candle._cython._dtype", None) + monkeypatch.delitem(sys.modules, "candle._dtype", raising=False) + + with pytest.raises(ImportError) as exc: + importlib.import_module("candle._dtype") + + msg = str(exc.value) + assert "_dtype" in msg + assert any( + hint in msg + for hint in ("build_ext", "pip install", "setup.py", "[cython]") + ) + + + +def test_cython_package_import_without_dispatcher_core_is_hard_failure(monkeypatch): + monkeypatch.setitem(sys.modules, "candle._cython._dispatcher_core", None) + monkeypatch.delitem(sys.modules, "candle._cython", raising=False) + + with pytest.raises(ImportError) as exc: + importlib.import_module("candle._cython") + + msg = str(exc.value) + assert "_dispatcher_core" in msg + assert any( + hint in msg + for hint in ("build_ext", "pip install", "setup.py", "[cython]") + ) + + + +def test_public_dtype_singletons_are_fastdtype_instances(): + import candle as torch + from candle._cython._dtype import FastDType + + assert isinstance(torch.float32, FastDType) + assert isinstance(torch.float16, FastDType) + assert isinstance(torch.int64, FastDType) + + + +def test_cython_package_runtime_core_flags_are_true(): + import candle._cython as c + + assert c._HAS_CYTHON_TENSOR_IMPL is True + assert c._HAS_CYTHON_DISPATCHER_CORE is True + assert c._HAS_CYTHON_DEVICE is True + assert c._HAS_CYTHON_DTYPE is True + assert c._HAS_CYTHON_DISPATCH is True + assert c._HAS_CYTHON_STORAGE is True + assert c._HAS_CYTHON_FAST_OPS is True + + + +def test_import_candle_succeeds_with_built_runtime_core(): + env = dict(os.environ) + env["PYTHONPATH"] = "src" + proc = subprocess.run( + [ + sys.executable, + "-c", + "import candle; import candle._dispatch.dispatcher as d; print('OK', d.dispatch.__module__)", + ], + cwd=os.path.dirname(os.path.dirname(os.path.dirname(__file__))), + env=env, + capture_output=True, + text=True, + check=False, + ) + + assert proc.returncode == 0, proc.stderr or proc.stdout + assert "_dispatcher_core" in proc.stdout diff --git a/tests/npu/test_npu_soc_policy.py b/tests/npu/test_npu_soc_policy.py index 9bfdeb40..151b22c9 100644 --- a/tests/npu/test_npu_soc_policy.py +++ b/tests/npu/test_npu_soc_policy.py @@ -102,6 +102,7 @@ def test_soc_310b_fallback_ops_cover_expected_watchlist_set(): "avg_pool2d", "adaptive_avg_pool2d", "einsum", + "allclose", } got = set(ops_soc.fallback_ops("310b")) assert got == expected