diff --git a/docs/changelog.rst b/docs/changelog.rst index 4e4d40f7..a05b1bf1 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -1,6 +1,17 @@ Changelog ========= +Version 0.6.2 +------------- +- Fixed crash in difference refinement with mismatched reflection files: HKL reindexing (``validate_hkl``/``remap``/``reduce_to_spacegroup``) now carries all per-reflection fields (including the anomalous bookkeeping read by ``hkl_for_sf()``) instead of a hardcoded subset. +- Added a metal shader for structure factor calculation +- Standardized structure factor calculation geometry +- Split gpu tests into cuda and mps +- Reworked VDW pair list creation +- Reworked backend dispatch: which kernel runs is now read from one declarative table per kernel family (device, dtype, availability probe, failure policy), replacing two hand-written if/elif ladders. +- Added preconditioned L-BFGS optimizer for joined refinement. + + Version 0.6.1 ------------- - Fixed bug in beta estimation that caused instability in GPU refinement diff --git a/logo/ASCII.txt b/logo/ASCII.txt new file mode 100644 index 00000000..132134bb --- /dev/null +++ b/logo/ASCII.txt @@ -0,0 +1,21 @@ + oo + oo + oo # + oo # + oo # + oo ## o + oo ##### oo + oo ###### oo + oo ##### oo + o ##====== o + oo =========== oo + oo ==#### oo + o ########## o + o ########## o + o #####== o + o #=========== o + oo ======== oo + oo ==== oo + ooo ooo + oooo oooo + oooooooooo \ No newline at end of file diff --git a/logo/torchref-a-ink-cage-red-helix-small.svg b/logo/torchref-a-ink-cage-red-helix-small.svg new file mode 100644 index 00000000..38413d52 --- /dev/null +++ b/logo/torchref-a-ink-cage-red-helix-small.svg @@ -0,0 +1,23 @@ + + TorchRef — small mark + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/logo/torchref-a-ink-cage-red-helix.svg b/logo/torchref-a-ink-cage-red-helix.svg new file mode 100644 index 00000000..4ffb8e04 --- /dev/null +++ b/logo/torchref-a-ink-cage-red-helix.svg @@ -0,0 +1,29 @@ + + TorchRef + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/logo/torchref-b-red-cage-ink-helix-small.svg b/logo/torchref-b-red-cage-ink-helix-small.svg new file mode 100644 index 00000000..4bbe6591 --- /dev/null +++ b/logo/torchref-b-red-cage-ink-helix-small.svg @@ -0,0 +1,23 @@ + + TorchRef — small mark + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/logo/torchref-b-red-cage-ink-helix.svg b/logo/torchref-b-red-cage-ink-helix.svg new file mode 100644 index 00000000..8877805f --- /dev/null +++ b/logo/torchref-b-red-cage-ink-helix.svg @@ -0,0 +1,29 @@ + + TorchRef + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/logo/torchref-c-all-red-small.svg b/logo/torchref-c-all-red-small.svg new file mode 100644 index 00000000..0426073e --- /dev/null +++ b/logo/torchref-c-all-red-small.svg @@ -0,0 +1,23 @@ + + TorchRef — small mark + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/logo/torchref-c-all-red.svg b/logo/torchref-c-all-red.svg new file mode 100644 index 00000000..242d09a3 --- /dev/null +++ b/logo/torchref-c-all-red.svg @@ -0,0 +1,29 @@ + + TorchRef + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/logo/torchref-d-mono-no-accent-small.svg b/logo/torchref-d-mono-no-accent-small.svg new file mode 100644 index 00000000..07d4bc46 --- /dev/null +++ b/logo/torchref-d-mono-no-accent-small.svg @@ -0,0 +1,23 @@ + + TorchRef — small mark + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/logo/torchref-d-mono-no-accent.svg b/logo/torchref-d-mono-no-accent.svg new file mode 100644 index 00000000..a6a679d1 --- /dev/null +++ b/logo/torchref-d-mono-no-accent.svg @@ -0,0 +1,29 @@ + + TorchRef + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/logo/torchref-e-knockout-on-ink-small.svg b/logo/torchref-e-knockout-on-ink-small.svg new file mode 100644 index 00000000..548c8bf9 --- /dev/null +++ b/logo/torchref-e-knockout-on-ink-small.svg @@ -0,0 +1,24 @@ + + TorchRef — small mark + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/logo/torchref-e-knockout-on-ink.svg b/logo/torchref-e-knockout-on-ink.svg new file mode 100644 index 00000000..b25f2e37 --- /dev/null +++ b/logo/torchref-e-knockout-on-ink.svg @@ -0,0 +1,30 @@ + + TorchRef + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/logo/torchref-f-on-red-field-small.svg b/logo/torchref-f-on-red-field-small.svg new file mode 100644 index 00000000..c5904dfa --- /dev/null +++ b/logo/torchref-f-on-red-field-small.svg @@ -0,0 +1,24 @@ + + TorchRef — small mark + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/logo/torchref-f-on-red-field.svg b/logo/torchref-f-on-red-field.svg new file mode 100644 index 00000000..216fe550 --- /dev/null +++ b/logo/torchref-f-on-red-field.svg @@ -0,0 +1,30 @@ + + TorchRef + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 1ecb1db1..7c032667 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "torchref" -version = "0.6.1" +version = "0.6.2" description = "Pytorch based crystallographic refinement" readme = "README.md" requires-python = ">=3.10" @@ -98,8 +98,35 @@ profile = "black" line_length = 88 [tool.pytest.ini_options] +# Used when pytest is invoked from the repo root. ``tests/pytest.ini`` takes +# over when invoked from inside ``tests/``; keep the two in step so behaviour +# does not depend on the working directory. Accelerator tests are gated by +# ``tests/conftest.py``, which runs whatever the host supports -- do not add a +# default ``-m`` expression here, as it deselects at collection time and no +# flag can undo it. testpaths = ["tests"] python_files = "test_*.py" +python_classes = "Test*" +python_functions = "test_*" +pythonpath = ["."] +markers = [ + "unit: Unit tests (fast, no I/O, isolated)", + "integration: Integration tests (slower, real I/O, pipelines)", + "gpu: Needs any accelerator (CUDA or MPS); auto-skipped if the host has none", + "cuda: Needs CUDA specifically (e.g. Triton kernels); auto-skipped if absent", + "mps: Needs MPS specifically (Metal kernels); auto-skipped if absent", + "cuda_only: Deprecated alias for 'cuda'", + "slow: Long-running tests (skipped in quick runs)", + "openmm: Tests requiring OpenMM (the [amber] extra); skipped if absent", + "amber: Tests requiring OpenMM + AmberTools (antechamber/tleap); skipped if absent", +] + +filterwarnings = [ + "ignore::DeprecationWarning", + "ignore::PendingDeprecationWarning", + "ignore::FutureWarning", + "error::torchref.utils.backends.TorchRefDegradationWarning", +] [tool.ruff.lint] ignore = ["F841","E741","F841","E402","E722"] \ No newline at end of file diff --git a/tests/README.md b/tests/README.md index 7b66ea92..fd283106 100644 --- a/tests/README.md +++ b/tests/README.md @@ -92,7 +92,10 @@ Tests are marked with the following pytest markers: - `@pytest.mark.unit` - Fast unit tests, no I/O - `@pytest.mark.integration` - Integration tests with file I/O -- `@pytest.mark.gpu` - Tests requiring GPU (not run by default) +- `@pytest.mark.gpu` - Needs any accelerator (CUDA *or* MPS); runs automatically + wherever one is present, skipped otherwise +- `@pytest.mark.cuda` - Needs CUDA specifically (e.g. Triton kernels) +- `@pytest.mark.mps` - Needs MPS specifically (Metal kernels) - `@pytest.mark.slow` - Slow tests (>30 seconds) ### Running by marker @@ -146,8 +149,16 @@ See `.github/workflows/tests.yml` for configuration. - Can use real files from `tests/files/` - Mark with `@pytest.mark.integration` -3. **GPU tests**: Add `@pytest.mark.gpu` marker - - Will be skipped if no GPU available +3. **Accelerator tests**: pick the marker that matches what the test actually + needs -- `@pytest.mark.cuda` if it hardcodes CUDA, `@pytest.mark.mps` if it + hardcodes MPS, `@pytest.mark.gpu` only if it works on either (take the device + from the `cuda_device` / `mps_device` / `gpu_device` fixture accordingly). + - The marker is the *only* gate. Do not also check availability inside the + test: a second `pytest.skip` can only mask a forgotten or wrong marker, + turning "this host cannot run it" into a silent pass. + - Missing backend -> skipped with a reason naming it. `--run-cuda` / + `--run-mps` instead warn and let the tests run, so they error with the real + backend error -- for CI that must not go green on a runner that lost its GPU. ## Test Data diff --git a/tests/conftest.py b/tests/conftest.py index 905a1953..84715a84 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,6 +5,7 @@ """ import importlib.util import shutil +import warnings import pytest import torchref @@ -21,13 +22,40 @@ _HAS_AMBERTOOLS = bool(shutil.which("antechamber") and shutil.which("tleap")) +def _cuda_available() -> bool: + return torch.cuda.is_available() + + +def _mps_available() -> bool: + return hasattr(torch.backends, "mps") and torch.backends.mps.is_available() + + def pytest_addoption(parser): """Add custom command line options.""" + parser.addoption( + "--run-cuda", + action="store_true", + default=False, + help=( + "Require CUDA: fail instead of skipping if it is unavailable. " + "CUDA tests already run automatically when a CUDA device is " + "present; use this in CI to catch a runner that lost its GPU." + ), + ) + parser.addoption( + "--run-mps", + action="store_true", + default=False, + help=( + "Require MPS: fail instead of skipping if it is unavailable. " + "MPS tests already run automatically on Apple silicon." + ), + ) parser.addoption( "--run-gpu", action="store_true", default=False, - help="Run GPU tests" + help="Deprecated no-op: accelerator tests now run automatically.", ) parser.addoption( "--run-slow", @@ -41,26 +69,101 @@ def pytest_configure(config): """Configure pytest markers.""" config.addinivalue_line("markers", "unit: Unit tests (fast, no I/O)") config.addinivalue_line("markers", "integration: Integration tests (slower, real I/O)") - config.addinivalue_line("markers", "gpu: GPU-requiring tests (CUDA or MPS; skipped by default)") - config.addinivalue_line("markers", "cuda_only: Tests that specifically require CUDA (e.g. Triton)") + config.addinivalue_line("markers", "gpu: Needs any accelerator (CUDA or MPS); auto-skipped if none") + config.addinivalue_line("markers", "cuda: Needs CUDA specifically (e.g. Triton); auto-skipped if absent") + config.addinivalue_line("markers", "mps: Needs MPS specifically (Metal kernels); auto-skipped if absent") + config.addinivalue_line("markers", "cuda_only: Deprecated alias for 'cuda'") config.addinivalue_line("markers", "slow: Slow tests (skipped by default)") config.addinivalue_line("markers", "openmm: Needs OpenMM (the [amber] extra); skipped if absent") config.addinivalue_line("markers", "amber: Needs OpenMM + AmberTools (antechamber/tleap); skipped if absent") + if config.getoption("--run-gpu"): + # UserWarning, not DeprecationWarning: pytest.ini filters the latter, + # and a silently-swallowed notice is worse than none when the whole + # point is telling someone their flag no longer does anything. + warnings.warn( + "--run-gpu is deprecated and does nothing: accelerator tests now " + "run automatically wherever the backend is available. Use " + "--run-cuda / --run-mps to *require* a backend (fail rather than " + "skip when it is missing).", + UserWarning, + stacklevel=2, + ) + def pytest_collection_modifyitems(config, items): - """Skip GPU/slow tests unless explicitly requested.""" - skip_gpu = pytest.mark.skip(reason="Need --run-gpu option to run") + """Gate tests on what this host can actually do. + + Accelerator tests are **not** opt-in: a ``cuda``-marked test runs whenever + CUDA is present, an ``mps``-marked test whenever MPS is, and a ``gpu``-marked + (backend-agnostic) test whenever either is. Anything the host cannot run is + skipped with a reason naming the missing backend. + + ``--run-cuda`` / ``--run-mps`` invert the *absence* case from skip to + error, for CI that expects a specific backend and would otherwise go green + on a runner that quietly lost its GPU. They do it by *not* adding the skip + marker, so the backend tests run and fail with the real error from torch. + + This function is the **only** place that decides what runs. Tests must not + re-check availability themselves: a second layer of ``pytest.skip`` can only + mask a forgotten marker, and turns "this host cannot run it" into a silent + pass instead of the visible skip or the real error. + """ + has_cuda = _cuda_available() + has_mps = _mps_available() + + # Forced-but-absent is a warning, not a ``pytest.UsageError``. A UsageError + # aborts the entire session -- every unrelated test with it -- and says + # nothing about which backend call actually broke. Warning and letting the + # marked tests run gives a precise per-test failure and still runs the rest + # of the suite. UserWarning, not DeprecationWarning: pytest.ini filters the + # latter (see the --run-gpu note in pytest_configure). + require_cuda = config.getoption("--run-cuda") + require_mps = config.getoption("--run-mps") + if require_cuda and not has_cuda: + warnings.warn( + "--run-cuda given but no CUDA device is available: running the " + "cuda-marked tests anyway so they error with the real backend " + "error instead of being skipped.", + UserWarning, + stacklevel=2, + ) + if require_mps and not has_mps: + warnings.warn( + "--run-mps given but MPS is not available: running the mps-marked " + "tests anyway so they error with the real backend error instead of " + "being skipped.", + UserWarning, + stacklevel=2, + ) + skip_slow = pytest.mark.skip(reason="Need --run-slow option to run") skip_openmm = pytest.mark.skip(reason="OpenMM not installed (pip install '.[amber]')") skip_amber = pytest.mark.skip( reason="AmberTools (antechamber/tleap) not on PATH (conda install ambertools)" ) + skip_cuda = pytest.mark.skip(reason="No CUDA device on this host") + skip_mps = pytest.mark.skip(reason="No MPS device on this host") + skip_gpu = pytest.mark.skip(reason="No accelerator (CUDA or MPS) on this host") for item in items: - if "gpu" in item.keywords and not config.getoption("--run-gpu"): + keywords = item.keywords + # ``cuda_only`` is the retired spelling of ``cuda``. + wants_cuda = "cuda" in keywords or "cuda_only" in keywords + wants_mps = "mps" in keywords + if wants_cuda and not has_cuda and not require_cuda: + item.add_marker(skip_cuda) + if wants_mps and not has_mps and not require_mps: + item.add_marker(skip_mps) + # A bare ``gpu`` mark means "any accelerator"; a test that also names a + # specific backend has already been gated on the stricter condition. + if ( + "gpu" in keywords + and not (wants_cuda or wants_mps) + and not (has_cuda or has_mps) + ): item.add_marker(skip_gpu) - if "slow" in item.keywords and not config.getoption("--run-slow"): + if "slow" in keywords and not config.getoption("--run-slow"): item.add_marker(skip_slow) # Amber stack gates: "amber" needs OpenMM + AmberTools; "openmm" needs # just OpenMM. Skip with the most specific missing-dependency reason. @@ -135,26 +238,120 @@ def cpu_device() -> torch.device: return torch.device("cpu") -def _gpu_available() -> bool: - """Return True if any GPU backend (CUDA or MPS) is available.""" - if torch.cuda.is_available(): - return True - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - return True - return False - - @pytest.fixture(scope="session") def gpu_device() -> torch.device: """GPU torch device (only use with @pytest.mark.gpu). - Prefers CUDA, falls back to MPS; skips if neither is available. + Prefers CUDA, falls back to MPS; skips if neither is available. Prefer the + backend-specific ``cuda_device`` / ``mps_device`` below when a test needs one + particular backend -- this fixture's preference order means a + ``cuda``-marked test asking for it on a dual-backend host could be handed + MPS, which is why the MPS tests used to carry a ``type != 'mps'`` skip to + undo it. + """ + accel = _accelerator() + if accel is None: + pytest.skip("No accelerator (CUDA or MPS) on this host") + return accel + + +@pytest.fixture(scope="session") +def cuda_device() -> torch.device: + """Canonical CUDA device for ``cuda``-marked tests. + + Deliberately unguarded. What runs is decided by the ``cuda`` marker in + :func:`pytest_collection_modifyitems` and nowhere else, so this fixture does + not re-check availability: on a host without CUDA the test is *meant* to + error with the real backend error rather than be quietly skipped here. """ - if torch.cuda.is_available(): - return torch.device("cuda") - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - return torch.device("mps") - pytest.skip("No GPU (CUDA or MPS) available") + return torch.device("cuda", 0) + + +@pytest.fixture(scope="session") +def mps_device() -> torch.device: + """Canonical MPS device for ``mps``-marked tests. + + Unguarded for the same reason as :func:`cuda_device` -- the ``mps`` marker + owns the decision. + """ + return torch.device("mps", 0) + + +def _accelerator() -> "torch.device | None": + """The canonical accelerator this host can actually use, or ``None``. + + Indices are filled in (``cuda:0`` / ``mps:0``) so the value compares equal + to a device read back off a real tensor -- ``torch.device('mps')`` and + ``torch.device('mps:0')`` are *not* equal even though they name the same + physical device. + """ + if _cuda_available(): + return torch.device("cuda", torch.cuda.current_device()) + if _mps_available(): + return torch.device("mps", 0) + return None + + +# Built at import time so the ``gpu`` mark is attached during *collection*. +# Adding it later (e.g. via ``request.node.add_marker`` inside the fixture) is +# too late for ``pytest_collection_modifyitems`` to gate on. +_DEVICE_PARAMS = [pytest.param(torch.device("cpu"), id="cpu")] +_ACCELERATOR = _accelerator() +if _ACCELERATOR is not None: + _DEVICE_PARAMS.append( + pytest.param( + _ACCELERATOR, + id=_ACCELERATOR.type, + # Backend-specific mark, so a CUDA-less host skips the cuda leg and + # a non-Mac skips the mps leg, each with an accurate reason. + marks=getattr(pytest.mark, _ACCELERATOR.type), + ) + ) + + +@pytest.fixture(params=_DEVICE_PARAMS) +def any_device(request) -> torch.device: + """Every device this host can actually use, one test run per device. + + The CPU leg always runs. The accelerator leg is ``gpu``-marked, so a plain + ``pytest`` run skips it and ``pytest --run-gpu`` picks up CUDA on a CUDA + box or MPS on a Mac. On a CPU-only host the accelerator parameter does not + exist at all, so there is no skip noise. + """ + return request.param + + +@pytest.fixture(scope="session") +def _device_model_cache() -> dict: + """``{device_str: ModelFT}`` built at most once per device, per session.""" + return {} + + +@pytest.fixture +def device_model_bundle(_device_model_cache, pdb_dir, any_device): + """A loaded model on ``any_device``, for target conformance tests. + + The existing ``loaded_model`` / ``model_and_data`` fixtures are + function-scoped and construct on the process default, so a + device-parametrized sweep over them would reload the structure once per + test per device. This caches one model per device instead. + + Shared mutable state: callers must treat the bundle as read-only. A test + that moves a *target* will drag the borrowed model with it, poisoning every + later test on that device -- see ``test_target_device_round_trip``, which + deliberately builds its own. + """ + key = str(any_device) + if key not in _device_model_cache: + pdb = pdb_dir / "1DAW.pdb" + if not pdb.exists(): + pytest.skip("1DAW.pdb fixture not present") + from torchref.model import ModelFT + + _device_model_cache[key] = ModelFT(device=any_device, verbose=0).load_pdb( + str(pdb) + ) + return {"model": _device_model_cache[key]} @pytest.fixture @@ -171,7 +368,7 @@ def device(request) -> torch.device: markers = {m.name for m in request.node.iter_markers()} if "cuda_only" in markers and not torch.cuda.is_available(): pytest.skip("Test requires CUDA") - if "gpu" in markers and not _gpu_available(): + if "gpu" in markers and not (_cuda_available() or _mps_available()): pytest.skip("No GPU (CUDA or MPS) available") return get_default_device() @@ -409,3 +606,34 @@ def model_with_restraints(loaded_model): ) restraints.build_restraints() return {"model": loaded_model, "restraints": restraints} + +@pytest.fixture +def double_cpu(): + """float64/complex128 on CPU for the duration of a test; restore afterwards. + + Required rather than cosmetic for anything touching eager structure factors: + ``iso_structure_factor_torched`` casts ``hkl`` to the *global* ``dtypes.float`` + (``torchref/base/direct_summation/isotropic.py:121``), so under the default float32 + config a float64 leaf produces a dtype-mismatched matmul. + + Promoted here from three byte-similar copies in ``tests/unit/test_kernel_fixes.py``, + ``tests/unit/test_gradient_correctness.py`` and + ``tests/integration/test_dtype_config_float64.py``. This version also restores + ``sigma_cutoff_ed``, which none of those did -- so a test that changed the cutoff + leaked it into everything that ran afterwards. + """ + import torchref + from torchref.config import device as _device, dtypes as _dtypes + + f0, c0, d0 = _dtypes.float, _dtypes.complex, _device.current + s0 = torchref.sigma_cutoff_ed.value + _dtypes.float = torch.float64 + _dtypes.complex = torch.complex128 + _device.current = torch.device("cpu") + try: + yield + finally: + _dtypes.float = f0 + _dtypes.complex = c0 + _device.current = d0 + torchref.sigma_cutoff_ed.value = s0 diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py new file mode 100644 index 00000000..52060c21 --- /dev/null +++ b/tests/helpers/__init__.py @@ -0,0 +1 @@ +"""Shared test helpers for the torchref suite.""" diff --git a/tests/helpers/device_asserts.py b/tests/helpers/device_asserts.py new file mode 100644 index 00000000..20fdd4bf --- /dev/null +++ b/tests/helpers/device_asserts.py @@ -0,0 +1,292 @@ +"""Device-consistency assertions for torchref objects. + +The walker here is written **independently** of +:func:`torchref.utils.device_mixin._apply_to_obj`. That is deliberate: an oracle +that mirrors the implementation it checks can only ever confirm the +implementation's own idea of what is reachable, so it would be blind to exactly +the tensors ``DeviceMixin`` forgets to move. This walker instead enumerates +everything it can reach by generic Python/PyTorch means -- ``__dict__``, +``__slots__``, module parameters/buffers/children, containers, dataclass +fields, and :class:`~torchref.utils.utils.ModuleReference` wrappers. + +Exports +------- +assert_device_consistent + Fail unless every reachable tensor *and* every device tracker sits on the + expected device. Reports all mismatches at once. +collect_device_map + The underlying survey, returned as ``{path: torch.device}``. Useful when a + test wants to assert something more specific than "all on one device". +HOST_SIDE + Paths that are intentionally CPU-resident regardless of the object's + device. Every entry needs a comment justifying it. +""" + +from __future__ import annotations + +import dataclasses +from typing import Any, Dict, Iterable, List, Optional, Set, Tuple + +import torch +from torch import nn + +from torchref.config import canonical_device + +__all__ = [ + "assert_device_consistent", + "collect_device_map", + "HOST_SIDE", +] + + +# Attribute *leaf names* that are legitimately CPU-resident even when the +# owning object lives on an accelerator. Keep this list short and justified -- +# it is the escape hatch that makes the conformance suite meaningful rather +# than a rubber stamp. +# +# NOTE: a registered ``nn.Module`` buffer can never be host-side in practice, +# because ``module.to(device)`` relocates it regardless of what this set says. +# Anything that genuinely must stay on the host belongs in plain Python state +# (or ``get_extra_state``), not in a buffer. If you find yourself wanting to add +# a buffer name here, convert the buffer instead. +HOST_SIDE: Set[str] = set() + + +# Attributes on ``nn.Module`` that hold PyTorch's own bookkeeping. They are +# traversed explicitly via ``named_parameters``/``named_buffers``/ +# ``named_children``, so walking the raw dicts as well would only duplicate work +# and produce confusing double paths. +_MODULE_INTERNALS = frozenset( + { + "_parameters", + "_buffers", + "_modules", + "_non_persistent_buffers_set", + "_backward_hooks", + "_backward_pre_hooks", + "_forward_hooks", + "_forward_hooks_with_kwargs", + "_forward_pre_hooks", + "_forward_pre_hooks_with_kwargs", + "_state_dict_hooks", + "_state_dict_pre_hooks", + "_load_state_dict_pre_hooks", + "_load_state_dict_post_hooks", + "_is_full_backward_hook", + "training", + } +) + +_TRACKER_ATTRS = ("device", "_device") +_DTYPE_TRACKER_ATTRS = ("dtype_float", "_dtype") + +_MAX_DEPTH = 12 + + +def _slot_names(obj: Any) -> Iterable[str]: + """Every ``__slots__`` entry declared anywhere in ``type(obj)``'s MRO.""" + for klass in type(obj).__mro__: + for name in getattr(klass, "__slots__", ()) or (): + yield name + + +def _iter_attributes(obj: Any) -> Iterable[Tuple[str, Any]]: + """Yield ``(name, value)`` for the object's own state. + + Covers ``__dict__`` (skipping ``nn.Module`` bookkeeping), ``__slots__``, and + dataclass fields -- a dataclass with ``__slots__`` exposes neither through + ``__dict__``, which is how ``_ReflectionSubset``-style views hide tensors + from a naive walk. + """ + seen: Set[str] = set() + + for name, value in list(getattr(obj, "__dict__", {}).items()): + if name in _MODULE_INTERNALS: + continue + seen.add(name) + yield name, value + + for name in _slot_names(obj): + if name in seen or name.startswith("__"): + continue + seen.add(name) + try: + yield name, getattr(obj, name) + except AttributeError: + continue # declared slot, never assigned + + if dataclasses.is_dataclass(obj) and not isinstance(obj, type): + for field in dataclasses.fields(obj): + if field.name in seen: + continue + seen.add(field.name) + try: + yield field.name, getattr(obj, field.name) + except AttributeError: + continue + + +def _walk( + obj: Any, + path: str, + found: Dict[str, torch.device], + trackers: Dict[str, torch.device], + dtypes: Dict[str, torch.dtype], + dtype_trackers: Dict[str, torch.dtype], + visited: Set[int], + depth: int, +) -> None: + if obj is None or depth > _MAX_DEPTH: + return + if isinstance(obj, (str, bytes, int, float, bool, complex, torch.dtype, torch.device)): + return + + if isinstance(obj, torch.Tensor): + found[path] = obj.device + # Only floating/complex tensors carry the configured float dtype; + # integer index buffers and boolean masks have their own and must not + # be swept into the comparison. + if obj.is_floating_point() or obj.is_complex(): + dtypes[path] = obj.dtype + return + + if id(obj) in visited: + return + visited.add(id(obj)) + + args = (found, trackers, dtypes, dtype_trackers, visited) + + if isinstance(obj, nn.Module): + for name, tensor in list(obj.named_parameters(recurse=False)) + list( + obj.named_buffers(recurse=False) + ): + if tensor is None: + continue + found[f"{path}.{name}"] = tensor.device + if tensor.is_floating_point() or tensor.is_complex(): + dtypes[f"{path}.{name}"] = tensor.dtype + for name, child in obj.named_children(): + _walk(child, f"{path}.{name}", *args, depth + 1) + + # ``ModuleReference`` deliberately hides its payload from PyTorch's module + # tree, so a parameters/buffers sweep never sees it -- but the referenced + # object still has to agree on a device with the referrer. + wrapped = getattr(obj, "_wrapped_module", None) + if wrapped is not None and not isinstance(obj, nn.Module): + _walk(wrapped, f"{path}->ref", *args, depth + 1) + return + + for name, value in _iter_attributes(obj): + if name in _TRACKER_ATTRS and isinstance(value, (torch.device, str)): + try: + trackers[f"{path}.{name}"] = torch.device(value) + except (RuntimeError, TypeError): + pass # a ``device``-named attribute that isn't a device + continue + if name in _DTYPE_TRACKER_ATTRS and isinstance(value, torch.dtype): + dtype_trackers[f"{path}.{name}"] = value + continue + if name in HOST_SIDE: + continue + _walk(value, f"{path}.{name}", *args, depth + 1) + + if isinstance(obj, dict): + for key, value in list(obj.items()): + _walk(value, f"{path}[{key!r}]", *args, depth + 1) + elif isinstance(obj, (list, tuple, set, frozenset)): + for index, value in enumerate(obj): + _walk(value, f"{path}[{index}]", *args, depth + 1) + + +def collect_device_map(obj: Any, name: str = "obj") -> Tuple[Dict, Dict, Dict, Dict]: + """Survey ``obj`` along both the device and dtype axes. + + Returns ``(tensor_devices, tracker_devices, tensor_dtypes, tracker_dtypes)``, + each ``{path: value}``. Paths are dotted attribute chains, with ``[i]`` / + ``['k']`` for container elements and ``->ref`` where the walk crossed a + :class:`ModuleReference`. + + The dtype maps cover only floating/complex tensors and the + ``dtype_float`` / ``_dtype`` trackers, mirroring the mixin's rule that the + two axes are resolved independently -- an integer index buffer says nothing + about the configured float dtype. + """ + tensors: Dict[str, torch.device] = {} + trackers: Dict[str, torch.device] = {} + tensor_dtypes: Dict[str, torch.dtype] = {} + tracker_dtypes: Dict[str, torch.dtype] = {} + _walk(obj, name, tensors, trackers, tensor_dtypes, tracker_dtypes, set(), 0) + return tensors, trackers, tensor_dtypes, tracker_dtypes + + +def assert_device_consistent( + obj: Any, + expected: Any, + *, + name: str = "obj", + ignore: Optional[Iterable[str]] = None, + expected_dtype: Optional[torch.dtype] = None, +) -> None: + """Assert every reachable tensor and tracker sits on ``expected``. + + Parameters + ---------- + obj + Object to survey. + expected + Target device. Compared through + :func:`torchref.config.canonical_device` on both sides, so a bare + ``mps`` and an indexed ``mps:0`` are treated as equal -- the helper + tests placement, not spelling. + name + Root label used to build the reported paths. + ignore + Path substrings to exclude. Use for genuinely borrowed state, not to + paper over a split object. + expected_dtype + When given, also assert every floating/complex tensor and every + ``dtype_float`` / ``_dtype`` tracker matches. This is the only way to + catch the ``torch.tensor(float(x))`` class of bug, where a tunable is + hardcoded float32 under a float64 configuration. + + Raises + ------ + AssertionError + Listing **every** mismatch. A half-moved object usually has several, + and reporting only the first turns one debugging session into five. + """ + expected_dev = canonical_device(expected) + ignore = tuple(ignore or ()) + + tensors, trackers, tensor_dtypes, tracker_dtypes = collect_device_map(obj, name) + problems: List[str] = [] + + def skipped(path: str) -> bool: + return any(frag in path for frag in ignore) + + for path, dev in sorted(tensors.items()): + if not skipped(path) and canonical_device(dev) != expected_dev: + problems.append(f"{path}: tensor on {dev}, expected {expected_dev}") + + for path, dev in sorted(trackers.items()): + if not skipped(path) and canonical_device(dev) != expected_dev: + problems.append(f"{path}: tracker == {dev}, expected {expected_dev}") + + if expected_dtype is not None: + for path, dt in sorted(tensor_dtypes.items()): + if not skipped(path) and dt != expected_dtype: + problems.append(f"{path}: tensor dtype {dt}, expected {expected_dtype}") + for path, dt in sorted(tracker_dtypes.items()): + if not skipped(path) and dt != expected_dtype: + problems.append( + f"{path}: dtype tracker == {dt}, expected {expected_dtype}" + ) + + if problems: + target = f"{expected_dev}" + if expected_dtype is not None: + target += f" / {expected_dtype}" + raise AssertionError( + f"{type(obj).__name__} is not consistently on {target} " + f"({len(problems)} mismatch(es)):\n " + "\n ".join(problems) + ) diff --git a/tests/helpers/device_cases.py b/tests/helpers/device_cases.py new file mode 100644 index 00000000..321a982b --- /dev/null +++ b/tests/helpers/device_cases.py @@ -0,0 +1,337 @@ +"""Registry of cheaply-constructible device-bearing objects for conformance tests. + +Each :class:`DeviceCase` builds one object on a requested device. The suite in +``tests/unit/test_device_conformance.py`` then asserts that everything the +object owns actually landed there, and that it survives a device round trip. + +``UNCOVERED`` is the other half of the contract: every device-bearing class in +the source tree must appear either here or there, with a reason. That is what +keeps the registry from quietly rotting as the package grows -- see +``test_every_device_bearing_class_is_accounted_for``. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List + +import torch + +__all__ = [ + "DeviceCase", + "CASES", + "TargetDeviceCase", + "TARGET_CASES", + "UNCOVERED", + "UNCOVERED_PREFIXES", +] + +_CELL = [50.0, 60.0, 70.0, 90.0, 90.0, 90.0] +_SG = "P 21 21 21" + + +@dataclass(frozen=True) +class DeviceCase: + """One constructible object under device test. + + Parameters + ---------- + name + Test id. + build + ``device -> object``. Must place everything on ``device``. + cls_name + Class this case covers, for the coverage guard. + tensor_free + True when the object legitimately owns no tensors (an empty shell). + Such objects can only be checked via their tracker. + ignore + Path fragments to exclude from the walk (borrowed state). + """ + + name: str + build: Callable[[torch.device], Any] + cls_name: str + tensor_free: bool = False + ignore: tuple = field(default_factory=tuple) + + +def _cell(device): + from torchref.symmetry import Cell + + return Cell(_CELL, device=device) + + +CASES: List[DeviceCase] = [ + DeviceCase("Cell", _cell, "Cell"), + DeviceCase( + "SpaceGroup", + lambda d: __import__( + "torchref.symmetry", fromlist=["SpaceGroup"] + ).SpaceGroup(_SG, device=d), + "SpaceGroup", + ), + # D4: device implied by the cell; the SpaceGroup used to be built from the + # raw (None) device argument and land on the process default instead. + DeviceCase( + "SfFFT_from_cell", + lambda d: __import__( + "torchref.model.sf_fft", fromlist=["SfFFT"] + ).SfFFT(cell=_cell(d), spacegroup=_SG, max_res=2.0), + "SfFFT", + ), + # D4: explicit device disagreeing with the supplied cell. + DeviceCase( + "SfFFT_explicit_device", + lambda d: __import__( + "torchref.model.sf_fft", fromlist=["SfFFT"] + ).SfFFT(cell=_cell("cpu"), spacegroup=_SG, max_res=2.0, device=d), + "SfFFT", + ), + DeviceCase( + "SfDS_from_cell", + lambda d: __import__( + "torchref.model.sf_ds", fromlist=["SfDS"] + ).SfDS(cell=_cell(d), spacegroup=_SG), + "SfDS", + ), + # D1: tensor-free shells, whose tracker is the only thing to check. + DeviceCase( + "ScalerBase_empty", + lambda d: __import__( + "torchref.scaling", fromlist=["ScalerBase"] + ).ScalerBase(device=d), + "ScalerBase", + tensor_free=True, + ), + DeviceCase( + "LossState", + lambda d: __import__( + "torchref.refinement.loss_state", fromlist=["LossState"] + ).LossState(device=d), + "LossState", + tensor_free=True, + ), + DeviceCase( + "SolventModel", + lambda d: __import__( + "torchref.scaling.solvent", fromlist=["SolventModel"] + ).SolventModel(device=d), + "SolventModel", + ), + # D9: empty-init wrappers used to drop the requested device entirely. + DeviceCase( + "MixedTensor_empty", + lambda d: __import__( + "torchref.model.parameter_wrappers", fromlist=["MixedTensor"] + ).MixedTensor(device=d), + "MixedTensor", + ), + DeviceCase( + "MixedTensor_populated", + lambda d: __import__( + "torchref.model.parameter_wrappers", fromlist=["MixedTensor"] + ).MixedTensor( + torch.arange(12, dtype=torch.float32).reshape(4, 3), + # A 2D value tensor takes a 1D per-row mask. + refinable_mask=torch.zeros(4, dtype=torch.bool), + device=d, + ), + "MixedTensor", + ), + DeviceCase( + "PositiveMixedTensor", + lambda d: __import__( + "torchref.model.parameter_wrappers", fromlist=["PositiveMixedTensor"] + ).PositiveMixedTensor( + torch.arange(1, 5, dtype=torch.float32), + refinable_mask=torch.zeros(4, dtype=torch.bool), + device=d, + ), + "PositiveMixedTensor", + ), + DeviceCase( + "OccupancyTensor_empty", + lambda d: __import__( + "torchref.model.parameter_wrappers", fromlist=["OccupancyTensor"] + ).OccupancyTensor(device=d), + "OccupancyTensor", + ), + DeviceCase( + "RigidXYZTensor_empty", + lambda d: __import__( + "torchref.model.rigid_xyz", fromlist=["RigidXYZTensor"] + ).RigidXYZTensor(device=d), + "RigidXYZTensor", + ), + DeviceCase( + "ReciprocalSymmetryGrid", + lambda d: __import__( + "torchref.symmetry", fromlist=["ReciprocalSymmetryGrid"] + ).ReciprocalSymmetryGrid(_SG, grid_shape=(16, 16, 16), device=d), + "ReciprocalSymmetryGrid", + ), + DeviceCase( + "MapSymmetryDirect", + lambda d: __import__( + "torchref.symmetry", fromlist=["MapSymmetryDirect"] + ).MapSymmetryDirect( + _SG, map_shape=(16, 16, 16), cell_params=_CELL, device=d + ), + "MapSymmetryDirect", + ), + DeviceCase( + "TensorMasks", + lambda d: __import__( + "torchref.utils", fromlist=["TensorMasks"] + ).TensorMasks(device=d), + "TensorMasks", + tensor_free=True, + ), + DeviceCase( + "ReflectionData_empty", + lambda d: __import__( + "torchref.io", fromlist=["ReflectionData"] + ).ReflectionData(device=d), + "ReflectionData", + ), + DeviceCase( + "ManualWeighting", + lambda d: __import__( + "torchref.refinement.weighting", fromlist=["ManualWeighting"] + ).ManualWeighting({"geometry": 1.0}, device=d), + "ManualWeighting", + tensor_free=True, + ), +] + + +@dataclass(frozen=True) +class TargetDeviceCase: + """A refinement target, which needs a real model/data/scaler to construct. + + Separate from :class:`DeviceCase` because ``build`` takes a fixture-supplied + bundle as well as a device. The bundle is already on ``device``; ``build`` + must not cache or mutate it. + + ``ignore`` defaults to the wrapped objects: a target *borrows* its model and + data, and the conformance walk would otherwise re-check the entire structure + through every target that points at it. + """ + + name: str + build: Callable[[Dict[str, Any], torch.device], Any] + cls_name: str + needs: tuple = ("model",) + ignore: tuple = ("_model", "_data", "_scaler", "_model_dark", "_model_light") + + +TARGET_CASES: List[TargetDeviceCase] = [ + TargetDeviceCase( + "NonBondedTarget", + lambda b, d: __import__( + "torchref.refinement.targets.geometry.non_bonded", + fromlist=["NonBondedTarget"], + ).NonBondedTarget(b["model"]), + "NonBondedTarget", + ), + TargetDeviceCase( + "ADPSimilarityTarget", + lambda b, d: __import__( + "torchref.refinement.targets.adp.similarity", + fromlist=["ADPSimilarityTarget"], + ).ADPSimilarityTarget(b["model"]), + "ADPSimilarityTarget", + ), + TargetDeviceCase( + "ADPLocalityTarget", + lambda b, d: __import__( + "torchref.refinement.targets.adp.locality", fromlist=["ADPLocalityTarget"] + ).ADPLocalityTarget(b["model"]), + "ADPLocalityTarget", + ), + # Owns no tensors at all -- the case that exercises the request-driven + # tracker path rather than the owned-tensor path. + TargetDeviceCase( + "BondTarget", + lambda b, d: __import__( + "torchref.refinement.targets.geometry.bonds", fromlist=["BondTarget"] + ).BondTarget(b["model"]), + "BondTarget", + ), +] + + +# Device-bearing classes deliberately not in CASES. Every entry needs a reason; +# "hard to build" is a reason, "didn't get to it" is not. +UNCOVERED: Dict[str, str] = { + # --- abstract / mixin bases: never instantiated directly ----------------- + "Target": "abstract base; covered through its concrete subclasses", + "ModelTarget": "abstract base; needs a loaded model", + "DataTarget": "abstract base; needs model + data + scaler", + "XrayTarget": "abstract base; needs model + data + scaler", + "GeometryTarget": "abstract base; needs a model with restraints", + "CrystalDataset": "abstract dataclass base; covered via ReflectionData", + "BaseWeighting": "abstract base; covered via ManualWeighting", + "Refinement": "abstract base; covered via LBFGSRefinement in integration", + "PassThroughTensor": "documented non-functional stub (parameter_wrappers.py)", + "ADPTarget": "abstract base; needs a model with ADPs", + "CombinedTargets": "composite container; needs its component targets", + "CombinedModelTargets": "composite container; needs a loaded model", + "CollectionXrayTarget": "abstract base; needs a dataset collection", + # --- need loaded structures/data: covered by integration tests ----------- + "Model": "needs a PDB; covered by tests/unit/model/test_model_state_dict_device.py", + "ModelFT": "needs a PDB; covered by tests/unit/model/test_model_state_dict_device.py", + "MixedModel": "needs two loaded models", + "ModelCollection": "needs several loaded models", + "_SharedMixedModel": "internal view owned by ModelCollection", + "Scaler": "needs a loaded model + data; covered in integration", + "RestraintsNew": "needs a model + monomer library", + "HydrogenTopology": "needs a built restraint topology", + "FrenchWilson": "needs loaded intensities", + "DatasetCollection": "needs several loaded datasets", + "FcalcDataset": "needs computed structure factors", + "Map": "needs data + model", + "DifferenceMap": "needs two datasets + a model", + "LBFGSRefinement": "full pipeline; covered in integration", + "MapSymmetry": "interpolation variant; needs a real map grid", + "CholeskyMixedTensor": "needs a valid ADP tensor; shares MixedTensor's paths", + "CollectionScaler": "needs a dataset collection", + "CollectionDifferenceTarget": "needs a dataset collection", + "CollectionMLTarget": "needs a dataset collection", + "CollectionRiceTarget": "needs a dataset collection", + "ADPEntropyTarget": "needs a model with ADPs", + "AngleTarget": "needs a model with restraints", + "ChiralTarget": "needs a model with restraints", + "BhattacharyyaXrayTarget": "needs a scaler with initialised bins; buffers audited", + "ReciprocalSymmetryExtractor": "needs an hkl tensor + SpaceGroup; device is " + "intentionally inherited from hkl, so the generic 'requested device' " + "assertion does not apply", + "CoordinateSimilarityTarget": "needs two loaded models", + "MultiModelADPTarget": "needs a model collection", + "MultiModelGeometryTarget": "needs a model collection", + "TotalGeometryTarget": "composite; needs a model with restraints", + "TotalADPTarget": "composite; needs a model with restraints", + "NonBondedHTarget": "needs a model with hydrogen topology", + "PlanarityTarget": "needs a model with restraints", + "RamachandranTarget": "needs a model with restraints", + "TorsionTarget": "needs a model with restraints", + "RigidBondTarget": "needs a model with restraints", + "ScalerLogScaleTrendTarget": "needs a scaler", + "ScalerURegularizationTarget": "needs a scaler", + "GaussianXrayTarget": "needs model + data + scaler", + "LeastSquaresXrayTarget": "needs model + data + scaler", + "MaximumLikelihoodXrayTarget": "needs model + data + scaler", + "RiceXrayTarget": "needs model + data + scaler", + "DifferenceXrayTarget": "needs two datasets", + "PhaseInformedDifferenceTarget": "needs two datasets + phases", + "RiceDifferenceTarget": "needs two datasets", + "TaylorCorrectedDifferenceTarget": "needs two datasets", + "RigidTransform": "alignment helper; needs a coordinate set", + "RigidBodyRefinement": "experimental; needs model + data", +} + +# Everything under torchref/experimental is out of scope for the conformance +# sweep: it is explicitly unstable, and several modules need optional +# dependencies (OpenMM, AmberTools, JAX) that are not installed in CI. +UNCOVERED_PREFIXES = ("torchref/experimental/",) diff --git a/tests/helpers/device_inventory.py b/tests/helpers/device_inventory.py new file mode 100644 index 00000000..a404ac8e --- /dev/null +++ b/tests/helpers/device_inventory.py @@ -0,0 +1,70 @@ +"""Static inventory of ``DeviceMixin`` subclasses in the torchref source tree. + +Why AST rather than ``DeviceMixin.__subclasses__()``: the runtime hook only +sees classes whose defining module has been imported, so a coverage guard built +on it silently shrinks whenever an import is removed or an optional dependency +is missing -- exactly when you most want to be told about a gap. Parsing the +source finds every class whether or not it is importable here. + +Inheritance is resolved **transitively**: ``PositiveMixedTensor(MixedTensor)`` +never names a mixin alias itself, but it is still a device-bearing class and +still needs conformance coverage. +""" + +from __future__ import annotations + +import ast +from pathlib import Path +from typing import Dict, Set, Tuple + +__all__ = ["MIXIN_ALIASES", "device_mixin_classes"] + +# Every spelling of the one mixin. ``DeviceMovementMixin`` and +# ``_NonModuleDeviceMixin`` are aliases kept for backward compatibility. +MIXIN_ALIASES = frozenset( + {"DeviceMixin", "DeviceMovementMixin", "_NonModuleDeviceMixin"} +) + + +def _base_names(node: ast.ClassDef) -> Set[str]: + """Base-class names for ``node``, ignoring subscripts and keywords.""" + names = set() + for base in node.bases: + if isinstance(base, ast.Name): + names.add(base.id) + elif isinstance(base, ast.Attribute): + names.add(base.attr) + return names + + +def device_mixin_classes(package_root: Path) -> Dict[str, str]: + """Return ``{class_name: "relative/path.py:lineno"}`` for device-bearing classes. + + A class qualifies if any of its bases is a mixin alias, or is another + qualifying class -- iterated to a fixed point so multi-level hierarchies + are captured. + """ + definitions: Dict[str, Tuple[Set[str], str]] = {} + + for path in sorted(package_root.rglob("*.py")): + try: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + except (SyntaxError, UnicodeDecodeError): # pragma: no cover + continue + rel = path.relative_to(package_root.parent) + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef): + definitions[node.name] = (_base_names(node), f"{rel}:{node.lineno}") + + qualifying: Dict[str, str] = {} + changed = True + while changed: + changed = False + for name, (bases, where) in definitions.items(): + if name in qualifying: + continue + if bases & MIXIN_ALIASES or bases & qualifying.keys(): + qualifying[name] = where + changed = True + + return qualifying diff --git a/tests/helpers/grad_asserts.py b/tests/helpers/grad_asserts.py new file mode 100644 index 00000000..9b9d877b --- /dev/null +++ b/tests/helpers/grad_asserts.py @@ -0,0 +1,135 @@ +"""Gradient-agreement assertions shared across kernel comparison tests. + +Comparing two kernels' gradients needs **two** metrics, not one. Cosine +similarity alone passes a systematic scale error -- a kernel that returns +``2 * grad`` is perfectly parallel to the reference -- so every assertion here +checks direction *and* magnitude. + +Extracted from ``tests/unit/test_gradient_correctness.py`` so that other suites could +reuse them instead of hand-rolling a bare cosine check. The original consumers were the +accelerator kernel tests; those are now ``tests/unit/structure_factor/``, which compares +every production kernel against a direct-summation oracle rather than against another +kernel. +""" + +from __future__ import annotations + +import torch + +__all__ = [ + "cosine_similarity", + "gradnorm_ratio", + "assert_grads_agree", + "got_ref_pairs", + "rel_error", + "central_fd_grad", + "hvp", + "hvp_central_fd", +] + + +def _flat_real(t: torch.Tensor) -> torch.Tensor: + """Flatten to a real 1-D CPU float64 vector, splitting complex into (re, im). + + The CPU hop is required, not cosmetic: these helpers accumulate in float64 + and MPS has no float64, so comparing accelerator gradients in place would + fail in the backend instead of in the assertion. + """ + t = t.detach().cpu() + if t.is_complex(): + return torch.view_as_real(t).reshape(-1).double() + return t.reshape(-1).double() + + +def cosine_similarity(a: torch.Tensor, b: torch.Tensor) -> float: + """Cosine similarity of two gradients (direction agreement).""" + return torch.nn.functional.cosine_similarity( + _flat_real(a), _flat_real(b), dim=0 + ).item() + + +def gradnorm_ratio(a: torch.Tensor, b: torch.Tensor) -> float: + """Ratio of gradient norms ``||a|| / ||b||`` (magnitude agreement).""" + nb = _flat_real(b).norm() + return (_flat_real(a).norm() / nb).item() + + +def got_ref_pairs(got, ref): + """Yield ``(name, got_grad, ref_grad)`` triples for a dict or sequence.""" + if isinstance(got, dict): + for k in got: + yield str(k), got[k], ref[k] + else: + for i, (g, r) in enumerate(zip(got, ref)): + yield str(i), g, r + + +def assert_grads_agree(got, ref, *, min_cos=0.9999, ratio_tol=1e-3, ctx=""): + """Assert per-leaf cosine ≈ 1 and gradnorm ratio ≈ 1 for matched grads.""" + for name, g, r in got_ref_pairs(got, ref): + cos = cosine_similarity(g, r) + ratio = gradnorm_ratio(g, r) + assert cos >= min_cos, f"{ctx}{name}: cosine {cos:.8f} < {min_cos}" + assert ( + abs(ratio - 1.0) <= ratio_tol + ), f"{ctx}{name}: gradnorm ratio {ratio:.8f} off 1 by > {ratio_tol}" + + +def rel_error(a: torch.Tensor, b: torch.Tensor) -> float: + """Relative L2 error ``||a - b|| / ||b||`` in float64, device-agnostic.""" + fa, fb = _flat_real(a), _flat_real(b) + return ((fa - fb).norm() / fb.norm()).item() + + +# --------------------------------------------------------------------------- +# Finite differences and Hessian-vector products +# --------------------------------------------------------------------------- +# Promoted here from three divergent private copies -- ``test_gradient_correctness.py`` +# (``_central_fd_grad``), ``test_triton_vs_eager_targets.py`` (``_hvp_vs_fd``) and +# ``test_kernel_fixes.py`` (inline, twice). They disagreed on eps and on whether to +# detach, which made their tolerances incomparable. +# +# Caveat on where these are valid: a central difference is only a sound reference for a +# smooth function. The electron-density map route is hard-truncated at a per-atom radius +# and the cull surface moves with the atom, so a voxel crossing it contributes a fixed +# jump to the difference. Use finite differences against direct summation, which is +# analytic, and use direct summation as the reference for the map route. +def central_fd_grad(loss_fn, x0: torch.Tensor, eps: float) -> torch.Tensor: + """Central-difference gradient of a scalar ``loss_fn(x)`` at ``x0``. + + Costs ``2 * x0.numel()`` forward evaluations, so keep ``x0`` small. + """ + flat = x0.detach().reshape(-1).clone() + out = torch.zeros_like(flat) + for i in range(flat.numel()): + plus, minus = flat.clone(), flat.clone() + plus[i] += eps + minus[i] -= eps + with torch.no_grad(): + hi = loss_fn(plus.view_as(x0)) + lo = loss_fn(minus.view_as(x0)) + out[i] = (hi - lo) / (2.0 * eps) + return out.view_as(x0) + + +def hvp(loss_fn, x0: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + """Hessian-vector product by double backward: ``d/dx (grad(loss) . v)``.""" + x = x0.detach().clone().requires_grad_(True) + (g1,) = torch.autograd.grad(loss_fn(x), x, create_graph=True) + (out,) = torch.autograd.grad((g1 * v).sum(), x) + return out.detach() + + +def hvp_central_fd(loss_fn, x0: torch.Tensor, v: torch.Tensor, eps: float): + """Hessian-vector product by central-differencing the *gradient* along ``v``. + + Two gradient evaluations rather than ``2 * numel`` forwards, so this stays cheap even + on realistic scenes. + """ + + def grad_at(xv): + x = xv.detach().clone().requires_grad_(True) + (g,) = torch.autograd.grad(loss_fn(x), x) + return g.detach() + + return (grad_at(x0 + eps * v) - grad_at(x0 - eps * v)) / (2.0 * eps) diff --git a/tests/integration/test_cli_difference_refine_mismatched_mtz.py b/tests/integration/test_cli_difference_refine_mismatched_mtz.py new file mode 100644 index 00000000..a1a2723b --- /dev/null +++ b/tests/integration/test_cli_difference_refine_mismatched_mtz.py @@ -0,0 +1,110 @@ +"""CPU end-to-end test for ``torchref.difference-refine`` with MISMATCHED MTZ files. + +Difference refinement builds a ``DatasetCollection`` from two independent MTZ +files and aligns the light dataset onto the dark reference via +``ReflectionData.validate_hkl``. When the two files had different reflection +sets, alignment left ``hkl_anomalous`` (read by ``hkl_for_sf``) at the +pre-alignment length and the run crashed inside the collection difference target +(``RuntimeError: The size of tensor a (...) must match the size of tensor b``). + +This test generates two MTZ files with genuinely different reflection sets from +the smallest fixture (3GR5) and asserts the CLI now completes. It is the +end-to-end guard for the 0.6.2 fix; pre-fix it exited non-zero. + +Wall-clock budget: a couple of minutes on CPU (1 macro-cycle, 1 step). +""" + +import json +import subprocess +import sys +from pathlib import Path + +import pytest + + +@pytest.fixture +def diff_refine_cli_script(project_root) -> Path: + script = project_root / "torchref" / "cli" / "collection_difference_refine.py" + if not script.exists(): + pytest.skip(f"difference-refine CLI script not found: {script}") + return script + + +@pytest.fixture +def mismatched_mtz_pair(test_files_dir, tmp_path): + """Two MTZ files sharing a cell/spacegroup but with different reflections. + + Built from 3GR5: ``dark`` drops the last 10% of reflections, ``light`` drops + the first 15%, so each contains reflections the other lacks and the two + counts differ -- exactly the condition that triggered the crash. + """ + pdb_file = test_files_dir / "pdb" / "3GR5.pdb" + mtz_file = test_files_dir / "mtz" / "3GR5.mtz" + if not pdb_file.exists() or not mtz_file.exists(): + pytest.skip("3GR5 test files not found") + + import torch + + from torchref.io.datasets.reflection_data import ReflectionData + + full = ReflectionData(device="cpu", verbose=0).load_mtz(str(mtz_file)) + n = len(full.hkl) + idx = torch.arange(n) + dark = full.__select__(idx < int(n * 0.90)) + light = full.__select__(idx >= int(n * 0.15)) + assert len(dark.hkl) != len(light.hkl), "subsets must differ in size" + + dark_mtz = tmp_path / "dark.mtz" + light_mtz = tmp_path / "light.mtz" + dark.write_mtz(str(dark_mtz)) + light.write_mtz(str(light_mtz)) + return {"pdb": pdb_file, "dark": dark_mtz, "light": light_mtz} + + +@pytest.mark.integration +def test_difference_refine_mismatched_mtz_cpu( + diff_refine_cli_script, mismatched_mtz_pair, tmp_path +): + outdir = tmp_path / "diff_out" + result = subprocess.run( + [ + sys.executable, + str(diff_refine_cli_script), + "-dm", str(mismatched_mtz_pair["pdb"]), + "-lm", str(mismatched_mtz_pair["pdb"]), + "-dsf", str(mismatched_mtz_pair["dark"]), + "-lsf", str(mismatched_mtz_pair["light"]), + "--fraction", "0.3", + "--n-cycles", "1", + "--n-steps", "1", + "--max-iter", "5", + "-o", str(outdir), + "--device", "cpu", + "--verbose", "1", + ], + capture_output=True, + text=True, + timeout=1800, + ) + + if result.returncode != 0: + print("STDOUT:", result.stdout[-3000:]) + print("STDERR:", result.stderr[-3000:]) + + assert result.returncode == 0, ( + f"difference-refine exited with {result.returncode} on mismatched MTZ " + f"files. stderr tail: {result.stderr[-800:]}" + ) + + # The exact stale-tensor shape mismatch must not reappear. + assert "must match the size of tensor" not in result.stderr + + prefix = "fractions_70_30" + summary = outdir / f"{prefix}_summary.json" + diff_mtz = outdir / f"{prefix}_difference_data.mtz" + assert summary.exists(), "summary JSON not written" + assert diff_mtz.exists(), "difference MTZ not written" + + with open(summary) as f: + data = json.load(f) + assert "results" in data and "r_factor_light" in data["results"] diff --git a/tests/integration/test_cli_refinement.py b/tests/integration/test_cli_refinement.py index d7f32783..c3325871 100644 --- a/tests/integration/test_cli_refinement.py +++ b/tests/integration/test_cli_refinement.py @@ -72,7 +72,7 @@ def test_cli_refine_missing_inputs(self, cli_script, tmp_path): @pytest.mark.integration @pytest.mark.slow - @pytest.mark.gpu + @pytest.mark.cuda def test_cli_refine_cuda_end_to_end(self, cli_script, h_structure_pair, tmp_path): """Run the refine CLI on CUDA end-to-end on an H-bearing structure. @@ -83,8 +83,6 @@ def test_cli_refine_cuda_end_to_end(self, cli_script, h_structure_pair, tmp_path previously crashed when VDW buffers were left on CPU after a mid-refinement rebuild (PR #19). """ - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") outdir = tmp_path / "refine_cuda" diff --git a/tests/integration/test_device_mixin.py b/tests/integration/test_device_mixin.py index 7e586442..b538e100 100644 --- a/tests/integration/test_device_mixin.py +++ b/tests/integration/test_device_mixin.py @@ -13,9 +13,6 @@ import pytest import torch -CUDA_AVAILABLE = torch.cuda.is_available() - - def _load_model_ft(pdb_file, mtz_file): """Helper: load a ModelFT and matching reflection data on CPU. @@ -36,6 +33,7 @@ def _load_model_ft(pdb_file, mtz_file): @pytest.mark.integration +@pytest.mark.cuda def test_modelft_cpu_gpu_cpu_sf_round_trip(sample_pdb_file, sample_mtz_file): """CPU -> GPU -> CPU structure-factor round-trip via the unified mixin. @@ -46,9 +44,6 @@ def test_modelft_cpu_gpu_cpu_sf_round_trip(sample_pdb_file, sample_mtz_file): 3. Move back to CPU, recompute Fcalc; verify CPU placement and values match the original CPU result. """ - if not CUDA_AVAILABLE: - pytest.skip("CUDA device not available") - model, data = _load_model_ft(sample_pdb_file, sample_mtz_file) hkl, *_ = data() @@ -66,17 +61,23 @@ def test_modelft_cpu_gpu_cpu_sf_round_trip(sample_pdb_file, sample_mtz_file): assert fcalc_cuda.device.type == "cuda" assert model.device.type == "cuda" assert model.cell.device.type == "cuda", "Cell did not migrate to GPU" - for buf in model.buffers(): - assert buf.device.type == "cuda", "buffer left behind on CPU" - break - - # Values must agree within numerical tolerance after the round-trip. The CPU - # (C++ box-separable splat) and GPU (Triton work-queue splat) paths differ - # slightly in truncation shape (box vs sphere), which surfaces mostly on the - # strongest low-resolution reflections in absolute terms. Aggregate agreement - # is excellent (~0.3% RMS, ~0.12% R-factor), so use a combined - # relative-OR-absolute tolerance rather than max-abs-vs-mean, which a single - # very strong reflection (|F| >> mean |F|) can trip on a tiny relative error. + for name, buf in model.named_buffers(): + assert buf.device.type == "cuda", f"buffer {name} left behind on CPU" + + # Values must agree within numerical tolerance after the round-trip. + # + # This tolerance is loose for a reason that no longer applies. It was set when the CPU + # path was a box-separable splat and the GPU path a sphere-truncated work-queue kernel, + # so the two genuinely differed in truncation shape (~0.3% RMS, ~0.12% R-factor, + # concentrated on the strongest low-resolution reflections). Every production kernel now + # applies the identical spherical cutoff -- see ``electron_density/main.py`` -- so the + # residual should be float32 kernel arithmetic only, which is far smaller. + # + # It is deliberately NOT tightened here: this test is CUDA-only and has not been run + # since the contract was unified, so any number picked now would be a guess. Re-measure + # on a GPU host and tighten then. The combined relative-OR-absolute form stays either + # way -- max-abs-vs-mean trips on a single very strong reflection (|F| >> mean |F|) for + # a tiny relative error. fcalc_cuda_cpu = fcalc_cuda.cpu() magnitude = fcalc_cpu_initial.abs().mean().item() max_abs_diff = (fcalc_cpu_initial - fcalc_cuda_cpu).abs().max().item() diff --git a/tests/integration/test_ds_triton_vs_eager.py b/tests/integration/test_ds_triton_vs_eager.py deleted file mode 100644 index 6e68a8a7..00000000 --- a/tests/integration/test_ds_triton_vs_eager.py +++ /dev/null @@ -1,129 +0,0 @@ -"""Triton-vs-eager equivalence for direct-summation structure factors. - -The custom Triton kernels (float32) are compared against the checkpointed -eager backend evaluated in float64 (downcast for comparison), for both the -isolated P1 kernels and the full ``SfDS`` symmetry loop. - -Markers: ``@pytest.mark.gpu`` (skipped without ``--run-gpu``) and -``@pytest.mark.integration``. Requires a CUDA device. -""" - -import pytest -import torch - -pytestmark = [pytest.mark.gpu, pytest.mark.integration] - -_ATOL = 1e-2 -_RTOL = 1e-3 - - -def _rel(a, b): - num = (a - b.to(a.dtype)).abs().max() - den = b.abs().max().to(num.dtype) + 1e-6 - return (num / den).item() - - -@pytest.fixture -def cuda(): - if not torch.cuda.is_available(): - pytest.skip("CUDA required for the Triton DS kernels") - return torch.device("cuda") - - -def _asym_loss(F, wr, wi): - return (wr * F.real).sum() + (wi * F.imag).sum() - - -def test_iso_triton_vs_eager(cuda): - from torchref.base.direct_summation.triton_ds import ds_iso_triton - from torchref.base.direct_summation import dispatch as D - - torch.manual_seed(0) - R, N = 400, 23 - hkl = torch.randint(-6, 7, (R, 3), device=cuda).float() - s = torch.rand(R, device=cuda) * 0.6 - A = torch.rand(N, 5, device=cuda) - B = torch.rand(N, 5, device=cuda) + 0.5 - wr = torch.randn(R, device=cuda) - wi = torch.randn(R, device=cuda) # asymmetric -> exposes a wrong imag sign - - xyz0 = torch.rand(N, 3, device=cuda, dtype=torch.float64) - occ0 = torch.rand(N, device=cuda, dtype=torch.float64) * 0.4 + 0.6 - adp0 = torch.rand(N, device=cuda, dtype=torch.float64) * 10 + 5 - - xt, ot, at = (t.float().detach().requires_grad_() for t in (xyz0, occ0, adp0)) - Ft = ds_iso_triton(hkl, s, xt, ot, at, A, B) - _asym_loss(Ft, wr, wi).backward() - - xr, orf, arf = (t.detach().requires_grad_() for t in (xyz0, occ0, adp0)) - Fr = D._checkpointed_iso(hkl, s, xr, orf, arf, A.double(), B.double(), max_memory_gb=None) - _asym_loss(Fr, wr.double(), wi.double()).backward() - - assert _rel(Ft, Fr) < _RTOL - assert _rel(xt.grad, xr.grad) < _RTOL - assert _rel(ot.grad, orf.grad) < _RTOL - assert _rel(at.grad, arf.grad) < _RTOL - - -def test_aniso_triton_vs_eager(cuda): - from torchref.base.direct_summation.triton_ds import ds_aniso_triton - from torchref.base.direct_summation import dispatch as D - - torch.manual_seed(1) - R, N = 400, 23 - hkl = torch.randint(-6, 7, (R, 3), device=cuda).float() - svec = torch.randn(R, 3, device=cuda) * 0.3 - A = torch.rand(N, 5, device=cuda) - B = torch.rand(N, 5, device=cuda) + 0.5 - wr = torch.randn(R, device=cuda) - wi = torch.randn(R, device=cuda) - - xyz0 = torch.rand(N, 3, device=cuda, dtype=torch.float64) - occ0 = torch.rand(N, device=cuda, dtype=torch.float64) * 0.4 + 0.6 - U0 = torch.rand(N, 6, device=cuda, dtype=torch.float64) * 0.04 + 0.01 - - xt, ot, Ut = (t.float().detach().requires_grad_() for t in (xyz0, occ0, U0)) - Ft = ds_aniso_triton(hkl, svec, xt, ot, Ut, A, B) - _asym_loss(Ft, wr, wi).backward() - - xr, orf, Ur = (t.detach().requires_grad_() for t in (xyz0, occ0, U0)) - Fr = D._checkpointed_aniso(hkl, svec, xr, orf, Ur, A.double(), B.double(), max_memory_gb=None) - _asym_loss(Fr, wr.double(), wi.double()).backward() - - assert _rel(Ft, Fr) < _RTOL - assert _rel(xt.grad, xr.grad) < _RTOL - assert _rel(ot.grad, orf.grad) < _RTOL - assert _rel(Ut.grad, Ur.grad) < _RTOL - - -def test_sfds_engine_toggle_end_to_end(cuda): - """Full SfDS symmetry loop: Engine.TRITON vs Engine.EAGER agree.""" - from torchref.symmetry import Cell - from torchref.model.sf_ds import SfDS - from torchref.base.direct_summation import Engine - - torch.manual_seed(2) - cell = Cell([50.0, 60.0, 70.0, 90.0, 90.0, 90.0]) - N, R = 30, 200 - hkl = torch.randint(-6, 7, (R, 3), device=cuda).float() - A = torch.rand(N, 5, device=cuda) - B = torch.rand(N, 5, device=cuda) + 0.5 - xyz0 = torch.rand(N, 3, device=cuda) * 20 - occ0 = torch.rand(N, device=cuda) * 0.4 + 0.6 - adp0 = torch.rand(N, device=cuda) * 10 + 5 - wr = torch.randn(R, device=cuda) - wi = torch.randn(R, device=cuda) - - def run(engine): - sf = SfDS(cell, spacegroup="P212121", engine=engine, max_memory_gb=2.0) - xyz = xyz0.clone().requires_grad_() - occ = occ0.clone().requires_grad_() - adp = adp0.clone().requires_grad_() - F, _ = sf.compute_structure_factors(hkl, xyz, adp, occ, A, B) - _asym_loss(F, wr, wi).backward() - return F.detach(), xyz.grad, occ.grad, adp.grad - - Ft = run(Engine.TRITON) - Fe = run(Engine.EAGER) - for t, e in zip(Ft, Fe): - assert _rel(t, e) < _RTOL diff --git a/tests/integration/test_dtype_config_float64.py b/tests/integration/test_dtype_config_float64.py index 02c029e6..942cca31 100644 --- a/tests/integration/test_dtype_config_float64.py +++ b/tests/integration/test_dtype_config_float64.py @@ -10,22 +10,6 @@ import pytest import torch -from torchref.config import device, dtypes - - -@pytest.fixture -def double_cpu(): - """float64/complex128 on CPU for the duration of a test; restore after.""" - f0, c0, d0 = dtypes.float, dtypes.complex, device.current - dtypes.float = torch.float64 - dtypes.complex = torch.complex128 - device.current = torch.device("cpu") - try: - yield - finally: - dtypes.float = f0 - dtypes.complex = c0 - device.current = d0 @pytest.mark.unit diff --git a/tests/integration/test_gpu_scale_stability.py b/tests/integration/test_gpu_scale_stability.py index 0cd5045d..a1f553ad 100644 --- a/tests/integration/test_gpu_scale_stability.py +++ b/tests/integration/test_gpu_scale_stability.py @@ -14,11 +14,9 @@ @pytest.mark.integration -@pytest.mark.gpu +@pytest.mark.cuda @pytest.mark.slow def test_gpu_5cycle_refinement_does_not_collapse_and_matches_cpu(pdb_dir, mtz_dir): - if not torch.cuda.is_available(): - pytest.skip("CUDA required") pdb = pdb_dir / "1DAW.pdb" mtz = mtz_dir / "1DAW.mtz" if not pdb.exists() or not mtz.exists(): diff --git a/tests/integration/test_refinement_pipeline.py b/tests/integration/test_refinement_pipeline.py index 78c1639a..4f385578 100644 --- a/tests/integration/test_refinement_pipeline.py +++ b/tests/integration/test_refinement_pipeline.py @@ -166,5 +166,5 @@ def test_gpu_refinement_setup(self, sample_structure_pair, gpu_device): data.load_mtz(sample_structure_pair["reflections"]) # Everything should be on GPU - assert model.xyz().device.type == 'cuda' - assert data.hkl.device.type == 'cuda' + assert model.xyz().device.type == gpu_device.type + assert data.hkl.device.type == gpu_device.type diff --git a/tests/integration/test_sfds_device.py b/tests/integration/test_sfds_device.py index e12d18ff..7a6ac700 100644 --- a/tests/integration/test_sfds_device.py +++ b/tests/integration/test_sfds_device.py @@ -38,12 +38,10 @@ def test_sfds_same_device_cpu(): assert torch.isfinite(F.real).all() and torch.isfinite(F.imag).all() -@pytest.mark.gpu +@pytest.mark.cuda @pytest.mark.integration def test_sfds_hkl_on_different_device(): """hkl on CPU while the module + atoms are on CUDA must not crash.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA required") cuda = torch.device("cuda") cell = Cell(_CELL, device=cuda) sf = SfDS(cell, spacegroup="P212121").to(cuda) diff --git a/tests/integration/test_triton_vs_eager_targets.py b/tests/integration/test_triton_vs_eager_targets.py index 3267118d..33e53068 100644 --- a/tests/integration/test_triton_vs_eager_targets.py +++ b/tests/integration/test_triton_vs_eager_targets.py @@ -9,10 +9,10 @@ and assert that loss values and parameter gradients agree to within a fp32 tolerance accounting for atomic-scatter non-determinism. -Toggling is done with the shared ``torchref.utils.use_engine`` context manager -(``Engine.EAGER`` vs ``Engine.TRITON``). No process restart needed. +Toggling is done with the shared ``torchref.utils.use_portable`` context manager +(``use_portable()`` vs the default). No process restart needed. -Markers: ``@pytest.mark.gpu`` (skipped by default) and +Markers: ``@pytest.mark.cuda`` (skipped by default) and ``@pytest.mark.integration``. Run on a CUDA box with ``pytest tests/integration/test_triton_vs_eager_targets.py -m gpu``. """ @@ -23,6 +23,8 @@ from pathlib import Path from typing import Dict, Tuple +import contextlib + import pytest import torch @@ -150,8 +152,6 @@ def _1daw_pair(pdb_dir, mtz_dir): @pytest.fixture(scope="module") def gpu_refinement(_1daw_pair): """Build a CUDA refinement once for the whole module.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") from torchref import LBFGSRefinement device = torch.device("cuda") @@ -177,7 +177,7 @@ def _target_names(state): return list(state.targets.keys()) -@pytest.mark.gpu +@pytest.mark.cuda @pytest.mark.integration @pytest.mark.parametrize( "target_name", @@ -191,34 +191,45 @@ def _target_names(state): "geometry/nonbonded", "geometry/ramachandran", "adp/simu", - "adp/locality", - "adp/KL", - "adp/scaler_U", - "adp/scaler_log_scale", ], ) def test_triton_matches_eager_per_target(target_name, gpu_refinement, gpu_state): - """For each registered target, forward + backward must match between - Triton and eager paths (up to the per-target tolerance).""" + """For each registered target, forward + backward must match between the Triton and + portable paths (up to the per-target tolerance). + + Four params were removed when this migrated off the ``Engine`` enum: + ``adp/locality``, ``adp/KL``, ``adp/scaler_U`` and ``adp/scaler_log_scale`` have **no + Triton implementation at all**, so they were comparing the eager path against itself and + passing at ``rel == 0``. That was true before the migration too -- forcing an engine had + never made them non-vacuous, because there was nothing on the other side. + + What keeps the remaining params honest is not the toggle. On a CUDA host + ``test_backend_is_available_where_it_is_expected`` fails if the Triton kernels *should* + work and do not, and a runtime fallback raises through the degradation warning. Either + fires before this test could quietly compare the reference against itself. + + Note the ``xray`` param is broader than a loss-math A/B: it runs through + ``LBFGSRefinement`` -> ``ModelFT`` -> ``build_electron_density``, so it also exercises the + density splat. That is deliberate coverage, not an accident of the harness. + """ if target_name not in gpu_state.targets: pytest.skip(f"target {target_name!r} not in this state") target = gpu_state.targets[target_name] - from torchref.utils import Engine, use_engine + from torchref.utils import use_portable - # ----- eager (Triton dispatch disabled) ----- - with use_engine(Engine.EAGER): + # ----- reference (dispatch pinned to the portable path) ----- + with use_portable(): eager = _run_target_capture_grads(target, gpu_refinement) - # ----- Triton (dispatch enabled) ----- - with use_engine(Engine.TRITON): - triton = _run_target_capture_grads(target, gpu_refinement) + # ----- Triton (the default on CUDA float32) ----- + triton = _run_target_capture_grads(target, gpu_refinement) atol, rtol = _tol_for(target_name) _assert_close(target_name, eager, triton, atol, rtol) -@pytest.mark.gpu +@pytest.mark.cuda @pytest.mark.integration @pytest.mark.parametrize("target_mode", ["bhattacharyya", "rice", "ls", "gaussian"]) def test_triton_matches_eager_xray_modes(target_mode, _1daw_pair): @@ -227,10 +238,8 @@ def test_triton_matches_eager_xray_modes(target_mode, _1daw_pair): a constructor argument. (The default 'ml' σ_A target is eager-only, so it has no Triton path to compare.) """ - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") from torchref import LBFGSRefinement - from torchref.utils import Engine, use_engine + from torchref.utils import use_portable device = torch.device("cuda") pdb, mtz = _1daw_pair @@ -245,9 +254,9 @@ def test_triton_matches_eager_xray_modes(target_mode, _1daw_pair): ref.scaler.refine_lbfgs() target = ref.xray_target_work - with use_engine(Engine.EAGER): + with use_portable(): eager = _run_target_capture_grads(target, ref) - with use_engine(Engine.TRITON): + with contextlib.nullcontext(): # the default: Triton on CUDA float32 triton = _run_target_capture_grads(target, ref) name = f"xray/{target_mode}" @@ -255,7 +264,7 @@ def test_triton_matches_eager_xray_modes(target_mode, _1daw_pair): _assert_close(name, eager, triton, atol, rtol) -@pytest.mark.gpu +@pytest.mark.cuda @pytest.mark.integration @pytest.mark.parametrize("n_atoms", [4, 5, 6]) def test_planarity_triton_per_atom_sigma(n_atoms): @@ -267,13 +276,11 @@ def test_planarity_triton_per_atom_sigma(n_atoms): non-uniform per-atom sigmas the old kernel diverged badly (>50% loss error). This pins eager == Triton for non-uniform sigmas. """ - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") from torchref.base.targets.planarity import ( _planarity_math_eager, planarity_math, ) - from torchref.utils import Engine, use_engine + from torchref.utils import use_portable dev = torch.device("cuda") g = torch.Generator(device=dev).manual_seed(7) @@ -285,10 +292,10 @@ def test_planarity_triton_per_atom_sigma(n_atoms): xe = base.clone().requires_grad_(True) xt = base.clone().requires_grad_(True) - with use_engine(Engine.EAGER): + with use_portable(): le = _planarity_math_eager(xe, [(idx, sigmas)]) (ge,) = torch.autograd.grad(le, xe) - with use_engine(Engine.TRITON): + with contextlib.nullcontext(): # the default: Triton on CUDA float32 lt = planarity_math(xt, [(idx, sigmas)]) (gt,) = torch.autograd.grad(lt, xt) @@ -298,7 +305,7 @@ def test_planarity_triton_per_atom_sigma(n_atoms): assert torch.allclose(ge, gt, atol=1e-4, rtol=1e-3) -@pytest.mark.gpu +@pytest.mark.cuda @pytest.mark.integration def test_geometry_degenerate_finite_grads(): """At degenerate geometry, BOTH eager and Triton give finite gradients. @@ -308,12 +315,10 @@ def test_geometry_degenerate_finite_grads(): singularities). Both backends are now floored with EPS so the gradient is finite (CPU == GPU behavior), avoiding NaN-poisoned refinement steps. """ - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") from torchref.base.targets.angle import angle_math from torchref.base.targets.bond import bond_math from torchref.base.targets.torsion import torsion_omega_math - from torchref.utils import Engine, use_engine + from torchref.utils import use_portable dev = torch.device("cuda") # Collinear chain of points -> degenerate angle & torsion; plus a @@ -352,9 +357,9 @@ def test_geometry_degenerate_finite_grads(): ), ] for name, fn, args in cases: - for engine in (Engine.EAGER, Engine.TRITON): + for pin in (True, False): x = xyz0.clone().requires_grad_(True) - with use_engine(engine): + with (use_portable() if pin else contextlib.nullcontext()): loss = fn(x, *args) (grad,) = torch.autograd.grad(loss, x) assert torch.isfinite(grad).all(), f"{name}/{engine}: non-finite grad" @@ -363,9 +368,9 @@ def test_geometry_degenerate_finite_grads(): def _hvp_vs_fd(model, hkl, eps=1e-5): """Return (cosine, rel_err) of the autograd Hessian-vector product vs a central finite difference of the gradient, through ``model.forward`` under - ``use_engine(Engine.EAGER)`` (the genuine pure-torch, double-differentiable + ``use_portable()`` (the genuine pure-torch, double-differentiable ED path).""" - from torchref.utils import Engine, use_engine + from torchref.utils import use_portable x = model.xyz.refinable_params x0 = x.detach().clone() @@ -377,7 +382,7 @@ def grad_at(xv): with torch.no_grad(): x.copy_(xv) model.reset_cache() - with use_engine(Engine.EAGER): + with use_portable(): sf = model(hkl, recalc=True) return ( torch.autograd.grad((sf.real**2 + sf.imag**2).sum(), x)[0].detach().clone() @@ -386,7 +391,7 @@ def grad_at(xv): with torch.no_grad(): x.copy_(x0) model.reset_cache() - with use_engine(Engine.EAGER): + with use_portable(): sf = model(hkl, recalc=True) (g1,) = torch.autograd.grad((sf.real**2 + sf.imag**2).sum(), x, create_graph=True) (hvp,) = torch.autograd.grad((g1 * v).sum(), x) @@ -403,19 +408,17 @@ def grad_at(xv): return cos, rel -@pytest.mark.gpu +@pytest.mark.cuda @pytest.mark.integration def test_eager_gpu_hessian_iso(tmp_path): - """`use_engine(Engine.EAGER)` gives correct GPU second derivatives (iso). + """``use_portable()`` gives correct GPU second derivatives (iso). The fast GPU density splat is first-order only (custom-Function backward with no grad_fn) and silently drops the second-order term under - create_graph. Under ``Engine.EAGER`` the iso splat now routes to the pure- + create_graph. Under ``use_portable()`` the iso splat now routes to the pure- torch (scatter_add) path, which composes under autograd. A Hessian-vector product through ``ModelFT.forward`` must then match finite differences. """ - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") import itertools from torchref.config import device, dtypes @@ -450,12 +453,10 @@ def test_eager_gpu_hessian_iso(tmp_path): dtypes.float, dtypes.complex = f0, c0 -@pytest.mark.gpu +@pytest.mark.cuda @pytest.mark.integration def test_eager_gpu_hessian_aniso(pdb_dir): """Same as above but for anisotropic ADPs (real ANISOU structure).""" - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") import itertools from torchref.config import dtypes diff --git a/tests/integration/test_variable_radius_gpu.py b/tests/integration/test_variable_radius_gpu.py deleted file mode 100644 index 23a2c07d..00000000 --- a/tests/integration/test_variable_radius_gpu.py +++ /dev/null @@ -1,196 +0,0 @@ -"""GPU variable-radius density path: Triton (wq_grid / wq_grid_aniso) vs the CPU -grouped splat, which share the identical per-atom radius policy. - -Both backends truncate each atom at its own ``N_sigma * sigma_eff`` radius, so the -forward maps must agree to float32 + analytic-vs-gathered-coord tolerance and the -gradients (xyz / adp / u / occ) must be parallel. Requires CUDA + Triton. - -Markers: ``@pytest.mark.gpu`` (skipped without ``--run-gpu``), ``integration``. -""" - -import math - -import pytest -import torch - -from torchref.base.electron_density.main import build_electron_density -from torchref.utils import Engine, use_engine - -pytestmark = [pytest.mark.gpu, pytest.mark.integration] - - -def _cell(): - a, b, c = 30.0, 25.0, 20.0 - nx, ny, nz = 60, 50, 40 - frac = torch.diag(torch.tensor([a, b, c])) - inv_frac = torch.diag(torch.tensor([1 / a, 1 / b, 1 / c])) - voxel = torch.tensor([a / nx, b / ny, c / nz]) - ii, jj, kk = torch.meshgrid( - torch.arange(nx), torch.arange(ny), torch.arange(nz), indexing="ij" - ) - rsg = (torch.stack([ii / nx, jj / ny, kk / nz], -1).to(torch.float32)) @ frac.T - return (a, b, c), (nx, ny, nz), frac.float(), inv_frac.float(), voxel.float(), rsg - - -def _cell_monoclinic(): - """Non-orthogonal (beta ~ 100 deg) cell to exercise the inv-frac-norm per-axis - box and the off-diagonal coordinate math (u_c gains an x-component).""" - a, b, c = 30.0, 25.0, 20.0 - nx, ny, nz = 60, 50, 40 - beta = math.radians(100.0) - frac = torch.tensor( - [[a, 0.0, c * math.cos(beta)], - [0.0, b, 0.0], - [0.0, 0.0, c * math.sin(beta)]], dtype=torch.float64) - inv_frac = torch.linalg.inv(frac) - voxel = (frac.norm(dim=0) / torch.tensor([nx, ny, nz], dtype=torch.float64)).float() - ii, jj, kk = torch.meshgrid( - torch.arange(nx), torch.arange(ny), torch.arange(nz), indexing="ij" - ) - fc = torch.stack([ii / nx, jj / ny, kk / nz], -1).double() - rsg = (fc @ frac.T).float() - return (a, b, c), (nx, ny, nz), frac.float(), inv_frac.float(), voxel, rsg - - -def _iso_atoms(cell, n=60, seed=0): - g = torch.Generator().manual_seed(seed) - a, b, c = cell - xyz = torch.rand(n, 3, generator=g) * torch.tensor([a, b, c]) - adp = torch.rand(n, generator=g) * 35 + 3 - occ = torch.ones(n) - A = torch.rand(n, 5, generator=g) * 5 - B = torch.rand(n, 5, generator=g) * 20 + 2 - return [t.float() for t in (xyz, adp, occ, A, B)] - - -def _aniso_atoms(cell, n=30, seed=1): - g = torch.Generator().manual_seed(seed) - a, b, c = cell - xyz = torch.rand(n, 3, generator=g) * torch.tensor([a, b, c]) - u = torch.zeros(n, 6) - u[:, :3] = torch.rand(n, 3, generator=g) * 0.12 + 0.02 - occ = torch.ones(n) - A = torch.rand(n, 5, generator=g) * 5 - B = torch.rand(n, 5, generator=g) * 20 + 2 - return [t.float() for t in (xyz, u, occ, A, B)] - - -def _cos(a, b): - a, b = a.reshape(-1).double(), b.reshape(-1).double() - return float((a @ b) / (a.norm() * b.norm() + 1e-12)) - - -def _rel_map(x, y): - return float((x - y).abs().max() / (y.abs().max() + 1e-8)) - - -def test_iso_triton_matches_cpu_grouped(gpu_device): - cell, grid, frac, inv_frac, voxel, rsg = _cell() - xyz, adp, occ, A, B = _iso_atoms(cell) - - with use_engine(Engine.AUTO): - cpu = build_electron_density( - rsg, xyz, adp, occ, A, B, inv_frac, frac, voxel - ) - with use_engine(Engine.TRITON): - gpu = build_electron_density( - rsg.to(gpu_device), xyz.to(gpu_device), adp.to(gpu_device), occ.to(gpu_device), - A.to(gpu_device), B.to(gpu_device), inv_frac.to(gpu_device), frac.to(gpu_device), voxel.to(gpu_device), - ) - assert _rel_map(gpu.cpu(), cpu) < 2e-2 - assert _cos(gpu.cpu(), cpu) > 0.9995 - - -def test_aniso_triton_matches_cpu_grouped(gpu_device): - cell, grid, frac, inv_frac, voxel, rsg = _cell() - xa, ua, oa, Aa, Ba = _aniso_atoms(cell) - empty = torch.zeros(0, 3) - - with use_engine(Engine.AUTO): - cpu = build_electron_density( - rsg, empty, torch.zeros(0), torch.zeros(0), torch.zeros(0, 5), - torch.zeros(0, 5), inv_frac, frac, voxel, - xyz_aniso=xa, u_aniso=ua, occ_aniso=oa, A_aniso=Aa, B_aniso=Ba, - ) - with use_engine(Engine.TRITON): - gpu = build_electron_density( - rsg.to(gpu_device), empty.to(gpu_device), torch.zeros(0, device=gpu_device), - torch.zeros(0, device=gpu_device), torch.zeros(0, 5, device=gpu_device), - torch.zeros(0, 5, device=gpu_device), inv_frac.to(gpu_device), frac.to(gpu_device), - voxel.to(gpu_device), - xyz_aniso=xa.to(gpu_device), u_aniso=ua.to(gpu_device), occ_aniso=oa.to(gpu_device), - A_aniso=Aa.to(gpu_device), B_aniso=Ba.to(gpu_device), - ) - assert _rel_map(gpu.cpu(), cpu) < 2e-2 - assert _cos(gpu.cpu(), cpu) > 0.9995 - - -def test_iso_triton_matches_eager_monoclinic(gpu_device): - """Guard the per-axis (inv-frac-norm) box + direct coords on a sheared cell: - TRITON (CUDA) vs the portable EAGER plain-scatter reference (independent box).""" - cell, grid, frac, inv_frac, voxel, rsg = _cell_monoclinic() - xyz, adp, occ, A, B = _iso_atoms(cell) - - with use_engine(Engine.EAGER): - ref = build_electron_density( - rsg, xyz, adp, occ, A, B, inv_frac, frac, voxel - ) - with use_engine(Engine.TRITON): - gpu = build_electron_density( - rsg.to(gpu_device), xyz.to(gpu_device), adp.to(gpu_device), occ.to(gpu_device), - A.to(gpu_device), B.to(gpu_device), inv_frac.to(gpu_device), frac.to(gpu_device), - voxel.to(gpu_device), - ) - assert _rel_map(gpu.cpu(), ref) < 2e-2 - assert _cos(gpu.cpu(), ref) > 0.9995 - - -def test_iso_gradients_match_cpu_grouped(gpu_device): - cell, grid, frac, inv_frac, voxel, rsg = _cell() - xyz, adp, occ, A, B = _iso_atoms(cell) - w = torch.randn(grid) - - def run(device, engine): - xx = xyz.to(device).clone().requires_grad_() - aa = adp.to(device).clone().requires_grad_() - oo = occ.to(device).clone().requires_grad_() - with use_engine(engine): - dm = build_electron_density( - rsg.to(device), xx, aa, oo, A.to(device), B.to(device), - inv_frac.to(device), frac.to(device), voxel.to(device), - ) - (dm * w.to(device)).sum().backward() - return xx.grad, aa.grad, oo.grad - - gx_c, ga_c, go_c = run("cpu", Engine.AUTO) - gx_g, ga_g, go_g = run(gpu_device, Engine.TRITON) - assert _cos(gx_g.cpu(), gx_c) > 0.999 - assert _cos(ga_g.cpu(), ga_c) > 0.999 - assert _cos(go_g.cpu(), go_c) > 0.999 - - -def test_aniso_gradients_match_cpu_grouped(gpu_device): - cell, grid, frac, inv_frac, voxel, rsg = _cell() - xa, ua, oa, Aa, Ba = _aniso_atoms(cell) - empty = torch.zeros(0, 3) - w = torch.randn(grid) - - def run(device, engine): - xx = xa.to(device).clone().requires_grad_() - uu = ua.to(device).clone().requires_grad_() - with use_engine(engine): - dm = build_electron_density( - rsg.to(device), empty.to(device), torch.zeros(0, device=device), - torch.zeros(0, device=device), torch.zeros(0, 5, device=device), - torch.zeros(0, 5, device=device), inv_frac.to(device), - frac.to(device), voxel.to(device), - xyz_aniso=xx, u_aniso=uu, occ_aniso=oa.to(device), - A_aniso=Aa.to(device), B_aniso=Ba.to(device), - ) - (dm * w.to(device)).sum().backward() - return xx.grad, uu.grad - - gx_c, gu_c = run("cpu", Engine.AUTO) - gx_g, gu_g = run(gpu_device, Engine.TRITON) - assert _cos(gx_g.cpu(), gx_c) > 0.999 - assert _cos(gu_g.cpu(), gu_c) > 0.999 diff --git a/tests/pytest.ini b/tests/pytest.ini index c8c990ef..46417e46 100644 --- a/tests/pytest.ini +++ b/tests/pytest.ini @@ -12,17 +12,25 @@ pythonpath = .. markers = unit: Unit tests (fast, no I/O, isolated) integration: Integration tests (slower, real I/O, pipelines) - gpu: Tests requiring CUDA GPU (skipped by default) + gpu: Needs any accelerator (CUDA or MPS); auto-skipped if the host has none + cuda: Needs CUDA specifically (e.g. Triton kernels); auto-skipped if absent + mps: Needs MPS specifically (Metal kernels); auto-skipped if absent + cuda_only: Deprecated alias for 'cuda' slow: Long-running tests (skipped in quick runs) openmm: Tests requiring OpenMM (the [amber] extra); skipped if absent amber: Tests requiring OpenMM + AmberTools (antechamber/tleap); skipped if absent -# Default options (exclude GPU tests by default) +# Default options. +# +# Accelerator tests are gated by ``pytest_collection_modifyitems`` in +# conftest.py, which runs whatever the host supports and skips the rest with a +# reason naming the missing backend. Do NOT add a default ``-m`` expression +# here: a mark expression *deselects* at collection time, which no flag can +# undo, so the gating would silently stop working. addopts = -v --tb=short --strict-markers - -m "not gpu" -ra # Ignore warnings from dependencies @@ -30,3 +38,10 @@ filterwarnings = ignore::DeprecationWarning ignore::PendingDeprecationWarning ignore::FutureWarning +# A backend degradation is an error under test and a warning in production. Dispatch falls +# back to the portable kernel when an accelerator throws, which is right for a user -- the +# answer stays correct, it just gets slower -- and wrong for CI, because a comparison between +# an accelerator and the reference would then run the reference twice and pass at zero +# difference while measuring nothing. Mirrored in ``pyproject.toml``, which is what applies +# when pytest is invoked from the repo root. + error::torchref.utils.backends.TorchRefDegradationWarning diff --git a/tests/scripts/submit_gpu_tests.sbatch b/tests/scripts/submit_gpu_tests.sbatch index 9d91951f..66435d92 100644 --- a/tests/scripts/submit_gpu_tests.sbatch +++ b/tests/scripts/submit_gpu_tests.sbatch @@ -50,9 +50,17 @@ echo "==========================================" echo "Running GPU tests..." echo "==========================================" -# Run GPU tests +# Run GPU tests. +# +# ``gpu or cuda`` rather than ``gpu`` alone: a mark expression *deselects* at collection +# time, so ``-m "gpu"`` never even collected the `cuda`-marked tests -- which is most of the +# accelerator coverage, including every direct-call kernel leg and the Triton provenance and +# gate tests. They were silently absent from this job rather than failing in it. +# +# ``mps`` is deliberately not included: this is a CUDA cluster, and those tests are skipped +# by ``pytest_collection_modifyitems`` on a host without Metal anyway. python -m pytest tests/ \ - -m "gpu" \ + -m "gpu or cuda" \ --cov=torchref \ --cov-report=term-missing \ -v \ diff --git a/tests/unit/base/test_aniso_map_building.py b/tests/unit/base/test_aniso_map_building.py deleted file mode 100644 index c343510d..00000000 --- a/tests/unit/base/test_aniso_map_building.py +++ /dev/null @@ -1,93 +0,0 @@ -"""Unit tests for anisotropic electron-density map building (CPU). - -Certifies the eager ``vectorized_add_to_map_aniso`` reference (iso-reduction -identity + float64 gradcheck). The production variable-radius CPU/GPU splats are -covered separately by ``test_variable_radius_gpu`` and the CPU grouped-splat tests. -""" - -import math - -import pytest -import torch - -from torchref.base.electron_density.map_building import vectorized_add_to_map_aniso -from torchref.base.electron_density.voxel_utils import find_relevant_voxels -from torchref.base.kernels import vectorized_add_to_map -from torchref.model.sf_fft import SfFFT -from torchref.symmetry import Cell - -pytestmark = pytest.mark.unit - -_EIGHT_PI_SQ = 8.0 * math.pi**2 - - -def _grid(dtype=torch.float64): - cell = Cell([12.0, 13.0, 14.0, 90.0, 90.0, 90.0]) - fft = SfFFT(cell, spacegroup="P1", max_res=3.0, - dtype_float=dtype, device="cpu") - fft.setup_grid(max_res=3.0) - return ( - fft.real_space_grid.to(dtype), - fft.inv_fractional_matrix.to(dtype), - fft.fractional_matrix.to(dtype), - 3.0, # splat radius for the legacy fixed-radius reference kernels - ) - - -def _atoms(dtype=torch.float64): - # fractional positions chosen away from voxel-center half-boundaries - xyz = torch.tensor([[4.1, 4.3, 5.2], [7.7, 8.9, 9.1]], dtype=dtype) - A = torch.rand(2, 5, dtype=dtype) + 0.5 - B = torch.rand(2, 5, dtype=dtype) + 1.0 - occ = torch.rand(2, dtype=dtype) + 0.5 - b = torch.tensor([8.0, 14.0], dtype=dtype) # isotropic B-factors - return xyz, A, B, occ, b - - -def test_aniso_reduces_to_isotropic(): - """U = b/(8π²)·I (diagonal, isotropic) must reproduce the iso splat.""" - grid, inv_frac, frac, rad = _grid() - xyz, A, B, occ, b = _atoms() - - surr, idx = find_relevant_voxels(grid, xyz, radius_angstrom=rad, inv_frac_matrix=inv_frac) - iso_map = vectorized_add_to_map( - surr, idx, torch.zeros(grid.shape[:3], dtype=torch.float64), - xyz, b, inv_frac, frac, A, B, occ, - ) - - u_iso = torch.zeros(2, 6, dtype=torch.float64) - u_iso[:, 0] = u_iso[:, 1] = u_iso[:, 2] = b / _EIGHT_PI_SQ - aniso_map = vectorized_add_to_map_aniso( - surr, idx, torch.zeros(grid.shape[:3], dtype=torch.float64), - xyz, u_iso, inv_frac, frac, A, B, occ, - ) - assert torch.allclose(iso_map, aniso_map, atol=1e-8, rtol=1e-6) - - -def test_aniso_map_gradcheck(): - grid, inv_frac, frac, rad = _grid() - xyz0, A, B, occ0, _ = _atoms() - # genuinely anisotropic U (positive-definite) - u0 = torch.tensor( - [[0.12, 0.10, 0.14, 0.02, 0.01, 0.015], - [0.09, 0.13, 0.11, -0.01, 0.02, 0.005]], - dtype=torch.float64, - ) - w = torch.randn(grid.shape[:3], dtype=torch.float64) - - xyz = xyz0.clone().requires_grad_() - u = u0.clone().requires_grad_() - occ = occ0.clone().requires_grad_() - - def f(xyz, u, occ): - surr, idx = find_relevant_voxels( - grid, xyz, radius_angstrom=rad, inv_frac_matrix=inv_frac - ) - dm = vectorized_add_to_map_aniso( - surr, idx, torch.zeros(grid.shape[:3], dtype=torch.float64), - xyz, u, inv_frac, frac, A, B, occ, - ) - return (dm * w).sum() - - assert torch.autograd.gradcheck(f, (xyz, u, occ), eps=1e-6, atol=1e-4) - diff --git a/tests/unit/base/test_canonical_sphere_cpu.py b/tests/unit/base/test_canonical_sphere_cpu.py new file mode 100644 index 00000000..943baac1 --- /dev/null +++ b/tests/unit/base/test_canonical_sphere_cpu.py @@ -0,0 +1,573 @@ +"""One truncation contract, checked on CPU with no accelerator required. + +Every production density kernel truncates each atom at the same spherical cutoff: + + voxel v gets atom i's density iff ||w||^2 <= r_i^2, with w the minimum-image + Cartesian atom->voxel vector (sphere centred on the ATOM, not on its anchor + node) and r_i the raw radius_policy radius, + +enumerated over the per-axis box ceil(r * n_axis * ||inv_frac row_axis||). This file +pins that contract down where it can actually run in CI on a laptop. + +Why this file exists +-------------------- +Every variable-radius cross-check used to be gated on hardware, so on a CPU-only +machine *nothing* tested the splat geometry. That is how three different cutoffs came to +coexist: the fast CPU path splatted a cube, the portable CPU path used a node-centred +diagonal metric at a voxel-rounded radius, and Metal inflated its radius to a whole +voxel. On a beta=115 deg cell those differed by ~5e-3 rel L2 -- more than the 1.7e-3 +truncation error the cutoff exists to deliver -- and the accelerator tests absorbed it in +a 2e-2 kernel-vs-kernel tolerance. + +Accuracy against an independent reference now lives in ``tests/unit/structure_factor/``, +which checks every production kernel on every device against direct summation. This file +remains the *geometry* contract: a from-spec brute force, cheap, CPU-only, and +independent of both kernels. + +The reference here is a direct O(atoms x voxels) evaluation of the contract as +written above (``_brute_iso`` / ``_brute_aniso``), not another kernel: comparing two +kernels to each other cannot catch a shared misreading of the geometry. + +Non-orthogonal cells are the point. An orthogonal cell hides a diagonal-metric +error entirely, so every geometry assertion runs at beta = 90, 100 and 115 degrees. +""" + +import contextlib +import math + +import pytest +import torch + +from torchref.base.electron_density.kernels.cpu import sphere_splat +from torchref.base.electron_density.kernels.cpu.variable_radius import ( + add_anisotropic_plain_var, + add_isotropic_plain_var, +) +from torchref.base.electron_density.main import build_electron_density +from torchref.base.electron_density.radius_policy import ( + per_atom_radius_aniso, + per_atom_radius_iso, +) +from torchref.base.scattering.scattering_table import get_scattering_params_by_z +from torchref.utils import use_portable + +pytestmark = pytest.mark.unit + +# float32 agreement floor. The fused kernel uses a fast 2^x exp (the CPU analogue of +# the metal::fast::exp the Metal kernels already use), measured at 2e-5 rel L2 +# against std::exp; the portable splat uses torch.exp. Both are far below the 7.9e-4 +# amplitude-truncation error at the default 3 sigma, so this tolerance bounds +# arithmetic noise while still failing on any geometry disagreement (the smallest of +# which, node- vs atom-centring on an orthogonal cell, is 9.9e-4). +_F32_TOL = 2e-4 +_BETAS = (90.0, 100.0, 115.0) + + +def _cell(beta_deg, dtype=torch.float32, dims=(48, 40, 34), abc=(28.0, 24.0, 20.0)): + """Monoclinic-family P1 cell; beta=90 degenerates to orthorhombic.""" + a, b, c = abc + beta = math.radians(beta_deg) + f64 = torch.tensor( + [[a, 0.0, c * math.cos(beta)], + [0.0, b, 0.0], + [0.0, 0.0, c * math.sin(beta)]], dtype=torch.float64) + return f64.to(dtype), torch.linalg.inv(f64).to(dtype), dims, f64 + + +def _voxel_size(f64, dims): + """``voxel_size`` as ``sf_fft`` derives it; unused by the kernels, still in the + ``build_electron_density`` signature.""" + return (f64.norm(dim=0) / torch.tensor(dims, dtype=torch.float64)).float() + + +def _iso_atoms(f64, n=36, dtype=torch.float32, seed=0): + g = torch.Generator().manual_seed(seed) + z = torch.tensor([6, 7, 8, 16]).repeat(n // 4 + 1)[:n] + A, B = get_scattering_params_by_z(z, dtype=dtype) + xyz = (torch.rand(n, 3, generator=g, dtype=torch.float64) @ f64.T).to(dtype) + adp = (torch.rand(n, generator=g) * 35 + 8).to(dtype) + # never exactly 1.0: the kernels recover d/d_occ by dividing the accumulated + # gradient by occ, and at occ == 1 a wrong scaling is invisible + occ = (torch.rand(n, generator=g) * 0.4 + 0.6).to(dtype) + return xyz, adp, occ, A, B + + +def _aniso_atoms(f64, n=24, dtype=torch.float32, seed=1): + g = torch.Generator().manual_seed(seed) + z = torch.tensor([6, 7, 8, 16]).repeat(n // 4 + 1)[:n] + A, B = get_scattering_params_by_z(z, dtype=dtype) + xyz = (torch.rand(n, 3, generator=g, dtype=torch.float64) @ f64.T).to(dtype) + u = torch.zeros(n, 6) + u[:, :3] = torch.rand(n, 3, generator=g) * 0.12 + 0.02 + # signed, non-zero off-diagonals: the p01/p02/p12 entries of the analytically + # inverted 3x3 and the 4*pi^2 off-diagonal U gradients are dead code without them + u[:, 3:] = (torch.rand(n, 3, generator=g) - 0.5) * 0.02 + occ = (torch.rand(n, generator=g) * 0.4 + 0.6).to(dtype) + return xyz, u.to(dtype), occ, A, B + + +def _frac_grid(dims, dtype): + """Voxel fractional coordinates, computed *in* ``dtype``. + + ``torch.arange(n) / n`` divides an integer tensor by a Python int, which lands in + the default float32 and would cap a float64 reference's accuracy at ~1e-7. + """ + ii, jj, kk = torch.meshgrid( + *[torch.arange(n, dtype=dtype) for n in dims], indexing="ij") + return torch.stack([ii / dims[0], jj / dims[1], kk / dims[2]], -1).reshape(-1, 3) + + +def _brute_iso(dims, xyz, adp, occ, A, B, inv_frac, frac, r): + """Direct evaluation of the canonical contract, atom by atom.""" + dtype = xyz.dtype + fc = _frac_grid(dims, dtype) + out = torch.zeros(dims[0] * dims[1] * dims[2], dtype=dtype) + xyz_frac = xyz @ inv_frac.T + B_total = ((B + adp[:, None]) * 0.25).clamp(min=0.1) + A_norm = A * occ[:, None] * (math.pi ** 1.5) / (B_total * torch.sqrt(B_total)) + for i in range(xyz.shape[0]): + d = fc - (xyz_frac[i] % 1.0) + d = d - torch.round(d) # minimum image + w = d @ frac.T + r2 = (w * w).sum(-1) + m = r2 <= r[i] * r[i] + e = torch.exp(-(math.pi ** 2) * r2[m][:, None] / B_total[i][None, :]) + out[m] = out[m] + (e * A_norm[i][None, :]).sum(-1) + return out.view(*dims) + + +def _brute_aniso(dims, xyz, u, occ, A, B, inv_frac, frac, r): + """As ``_brute_iso``, with the Mahalanobis density and a Euclidean cutoff.""" + dtype = xyz.dtype + fc = _frac_grid(dims, dtype) + out = torch.zeros(dims[0] * dims[1] * dims[2], dtype=dtype) + xyz_frac = xyz @ inv_frac.T + U3 = torch.zeros(u.shape[0], 3, 3, dtype=dtype) + U3[:, 0, 0], U3[:, 1, 1], U3[:, 2, 2] = u[:, 0], u[:, 1], u[:, 2] + U3[:, 0, 1] = U3[:, 1, 0] = u[:, 3] + U3[:, 0, 2] = U3[:, 2, 0] = u[:, 4] + U3[:, 1, 2] = U3[:, 2, 1] = u[:, 5] + M = (B[:, :, None, None] * torch.eye(3, dtype=dtype) + + 8 * math.pi ** 2 * U3[:, None]) / 4.0 + Minv = torch.linalg.inv(M) + A_norm = (A * occ[:, None] * (math.pi ** 1.5) + / torch.sqrt(torch.linalg.det(M).clamp(min=1e-10))) + for i in range(xyz.shape[0]): + d = fc - (xyz_frac[i] % 1.0) + d = d - torch.round(d) + w = d @ frac.T + m = (w * w).sum(-1) <= r[i] * r[i] # cutoff is the Euclidean sphere + ww = w[m] + q = torch.einsum("vi,gij,vj->vg", ww, Minv[i], ww) + out[m] = out[m] + (A_norm[i][None, :] * torch.exp(-(math.pi ** 2) * q)).sum(-1) + return out.view(*dims) + + +def _rel_l2(x, y): + return float((x - y).double().norm() / y.double().norm()) + + +def _empty_iso(dtype): + """Zero-length isotropic arguments, for an aniso-only structure.""" + return (torch.zeros(0, 3, dtype=dtype), torch.zeros(0, dtype=dtype), + torch.zeros(0, dtype=dtype), torch.zeros(0, 5, dtype=dtype), + torch.zeros(0, 5, dtype=dtype)) + + +# =========================================================================== +# The contract: each kernel vs a direct evaluation of the spec +# =========================================================================== + +@pytest.mark.parametrize("beta", _BETAS) +def test_fused_iso_matches_contract(beta): + if not sphere_splat.sphere_splat_available(): + pytest.skip(f"fused CPU splat unavailable: {sphere_splat.last_error()}") + frac, inv_frac, dims, f64 = _cell(beta) + xyz, adp, occ, A, B = _iso_atoms(f64) + r = per_atom_radius_iso(adp, B, n_sigma=3.0) + got = sphere_splat.add_isotropic_cpu_sphere_var( + torch.zeros(dims), xyz, adp, occ, A, B, inv_frac, frac, r) + assert _rel_l2(got, _brute_iso(dims, xyz, adp, occ, A, B, inv_frac, frac, r)) < _F32_TOL + + +@pytest.mark.parametrize("beta", _BETAS) +def test_fused_aniso_matches_contract(beta): + if not sphere_splat.sphere_splat_available(): + pytest.skip(f"fused CPU splat unavailable: {sphere_splat.last_error()}") + frac, inv_frac, dims, f64 = _cell(beta) + xyz, u, occ, A, B = _aniso_atoms(f64) + r = per_atom_radius_aniso(B, u, n_sigma=3.0) + got = sphere_splat.add_anisotropic_cpu_sphere_var( + torch.zeros(dims), xyz, u, occ, A, B, inv_frac, frac, r) + assert _rel_l2(got, _brute_aniso(dims, xyz, u, occ, A, B, inv_frac, frac, r)) < _F32_TOL + + +@pytest.mark.parametrize("beta", _BETAS) +def test_portable_iso_matches_contract(beta): + """The portable splat used a node-centred diagonal metric at a voxel-rounded + radius; on beta=115 that alone was 5e-3 rel L2.""" + frac, inv_frac, dims, f64 = _cell(beta) + xyz, adp, occ, A, B = _iso_atoms(f64) + r = per_atom_radius_iso(adp, B, n_sigma=3.0) + got = add_isotropic_plain_var( + torch.zeros(dims), xyz, adp, occ, A, B, inv_frac, frac, r) + assert _rel_l2(got, _brute_iso(dims, xyz, adp, occ, A, B, inv_frac, frac, r)) < _F32_TOL + + +@pytest.mark.parametrize("beta", _BETAS) +def test_portable_aniso_matches_contract(beta): + """The portable aniso splat used a full cube -- ~2.3x the sphere's voxels.""" + frac, inv_frac, dims, f64 = _cell(beta) + xyz, u, occ, A, B = _aniso_atoms(f64) + r = per_atom_radius_aniso(B, u, n_sigma=3.0) + got = add_anisotropic_plain_var( + torch.zeros(dims), xyz, u, occ, A, B, inv_frac, frac, r) + assert _rel_l2(got, _brute_aniso(dims, xyz, u, occ, A, B, inv_frac, frac, r)) < _F32_TOL + + +def test_fused_float64_is_exact(): + """float64 uses std::exp, so only fp rounding separates it from the reference.""" + if not sphere_splat.sphere_splat_available(): + pytest.skip("fused CPU splat unavailable") + frac, inv_frac, dims, f64 = _cell(100.0, dtype=torch.float64, dims=(32, 28, 24)) + xyz, adp, occ, A, B = _iso_atoms(f64, dtype=torch.float64) + r = per_atom_radius_iso(adp, B, n_sigma=3.0) + got = sphere_splat.add_isotropic_cpu_sphere_var( + torch.zeros(dims, dtype=torch.float64), xyz, adp, occ, A, B, inv_frac, frac, r) + want = _brute_iso(dims, xyz, adp, occ, A, B, inv_frac, frac, r) + assert _rel_l2(got, want) < 1e-13 + + +# =========================================================================== +# AUTO vs EAGER through the real dispatch: no accelerator needed +# =========================================================================== + +def _build(pin, dims, frac, inv_frac, voxel, dtype, iso=None, aniso=None): + rsg = torch.zeros(*dims, 3, dtype=dtype) # shape only; no kernel reads its values + xi, ai, oi, Ai, Bi = iso if iso is not None else _empty_iso(dtype) + kw = {} + if aniso is not None: + xa, ua, oa, Aa, Ba = aniso + kw = dict(xyz_aniso=xa, u_aniso=ua, occ_aniso=oa, A_aniso=Aa, B_aniso=Ba) + with (use_portable() if pin else contextlib.nullcontext()): + return build_electron_density( + rsg, xi, ai, oi, Ai, Bi, inv_frac, frac, voxel, dtype=dtype, **kw) + + +@pytest.mark.parametrize("beta", _BETAS) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64], + ids=["float32", "float64"]) +def test_auto_matches_eager_iso(beta, dtype): + frac, inv_frac, dims, f64 = _cell(beta, dtype=dtype) + voxel = _voxel_size(f64, dims) + atoms = _iso_atoms(f64, dtype=dtype) + ref = _build(True, dims, frac, inv_frac, voxel, dtype, iso=atoms) + got = _build(False, dims, frac, inv_frac, voxel, dtype, iso=atoms) + tol = _F32_TOL if dtype is torch.float32 else 1e-12 + assert _rel_l2(got, ref) < tol + + +@pytest.mark.parametrize("beta", _BETAS) +def test_auto_matches_eager_aniso(beta): + frac, inv_frac, dims, f64 = _cell(beta) + voxel = _voxel_size(f64, dims) + atoms = _aniso_atoms(f64) + ref = _build(True, dims, frac, inv_frac, voxel, torch.float32, aniso=atoms) + got = _build(False, dims, frac, inv_frac, voxel, torch.float32, aniso=atoms) + assert _rel_l2(got, ref) < _F32_TOL + + +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +def test_auto_matches_eager_gradients(kind): + """Direction *and* magnitude: a kernel returning ``2 * grad`` is perfectly + parallel, so cosine alone cannot catch it.""" + frac, inv_frac, dims, f64 = _cell(100.0, dtype=torch.float64, dims=(32, 28, 24)) + voxel = _voxel_size(f64, dims) + w = torch.randn(dims, generator=torch.Generator().manual_seed(7), + dtype=torch.float64) + if kind == "iso": + xyz, p, occ, A, B = _iso_atoms(f64, dtype=torch.float64) + else: + xyz, p, occ, A, B = _aniso_atoms(f64, dtype=torch.float64) + + def grads(pin): + x, pp, o = (t.clone().requires_grad_() for t in (xyz, p, occ)) + pack = (x, pp, o, A, B) + dm = _build(pin, dims, frac, inv_frac, voxel, torch.float64, + **({"iso": pack} if kind == "iso" else {"aniso": pack})) + (dm * w).sum().backward() + return x.grad, pp.grad, o.grad + + for g_auto, g_eager in zip(grads(False), grads(True)): + assert torch.allclose(g_auto, g_eager, rtol=1e-9, atol=0) + + +def test_cutoff_is_grid_independent(): + """The cutoff must not be requantized to the grid. + + Metal used to round the radius up to a whole voxel, which made the *same* + ``sigma_cutoff_ed`` truncate at different radii on different grid samplings. With + the radius used raw, refining the grid must converge the map rather than move the + cutoff: the total mass a coarse and a fine grid assign is set by the same sphere. + """ + frac, inv_frac, _, f64 = _cell(100.0) + xyz, adp, occ, A, B = _iso_atoms(f64, n=24) + r = per_atom_radius_iso(adp, B, n_sigma=3.0) + masses = [] + for dims in ((40, 34, 28), (60, 51, 42)): + dm = add_isotropic_plain_var( + torch.zeros(dims), xyz, adp, occ, A, B, inv_frac, frac, r) + cell_vol = float(torch.linalg.det(f64)) + masses.append(float(dm.double().sum()) * cell_vol / (dims[0] * dims[1] * dims[2])) + # Same truncation sphere, so the integrated mass agrees to grid-sampling error. + assert abs(masses[1] - masses[0]) / masses[0] < 5e-3 + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64], + ids=["float32", "float64"]) +def test_auto_actually_dispatches_the_fused_kernel(dtype, monkeypatch): + """Guard against the AUTO-vs-EAGER tests going vacuous. + + If the ``cpu_sphere`` row ever stopped matching, AUTO would fall through to the very + same portable splat EAGER uses and every equivalence assertion above would pass at + ``rel_l2 == 0`` while measuring nothing. float64 is parametrized because the dispatch + was once float32-only, so a regression to a dtype gate would be invisible on float32 + alone. + + The kernel's **defining module** is patched -- the same rule as every other provenance + test, since dispatch resolves each kernel by ``(module_path, attr)`` at call time. This + used to patch ``main`` instead, because the ladder there resolved the name from its own + globals. + """ + if not sphere_splat.sphere_splat_available(): + pytest.skip("fused CPU splat unavailable") + frac, inv_frac, dims, f64 = _cell(100.0, dtype=dtype) + calls = [] + real = sphere_splat.add_isotropic_cpu_sphere_var + + def recording(*args, **kwargs): + calls.append(1) + return real(*args, **kwargs) + + monkeypatch.setattr(sphere_splat, "add_isotropic_cpu_sphere_var", recording) + _build(False, dims, frac, inv_frac, _voxel_size(f64, dims), dtype, + iso=_iso_atoms(f64, dtype=dtype)) + assert calls, f"the default did not reach the fused CPU splat for {dtype}" + + +def test_empty_atom_sets(): + """A structure with no isotropic (or no anisotropic) atoms must not crash.""" + frac, inv_frac, dims, f64 = _cell(90.0) + voxel = _voxel_size(f64, dims) + only_aniso = _build(False, dims, frac, inv_frac, voxel, torch.float32, + aniso=_aniso_atoms(f64)) + assert torch.isfinite(only_aniso).all() and float(only_aniso.abs().sum()) > 0 + both_empty = _build(False, dims, frac, inv_frac, voxel, torch.float32) + assert float(both_empty.abs().sum()) == 0.0 + + +def test_density_map_accumulates_not_overwrites(): + """Both passes add into one map, so the aniso pass must not clobber the iso one.""" + frac, inv_frac, dims, f64 = _cell(90.0) + voxel = _voxel_size(f64, dims) + iso, aniso = _iso_atoms(f64), _aniso_atoms(f64) + a = _build(False, dims, frac, inv_frac, voxel, torch.float32, iso=iso) + b = _build(False, dims, frac, inv_frac, voxel, torch.float32, aniso=aniso) + both = _build(False, dims, frac, inv_frac, voxel, torch.float32, + iso=iso, aniso=aniso) + assert _rel_l2(both, a + b) < 1e-6 + + +# ========================================================================= +# Ported from now-deleted tests of orphaned kernels +# ========================================================================= +# ``tests/unit/base/test_aniso_map_building.py`` and ``test_cpu_scatter.py`` covered the +# legacy fixed-radius splat and the C++ structured scatter. The dispatch routes CPU to the +# fused sphere splat or the portable one, so that whole chain -- the grouped-separable and +# cube splats, the separable density core, the structured scatter -- became unreachable +# from ``build_electron_density`` and has since been deleted. The two checks worth keeping +# are re-pointed at the live path here. + + +@pytest.mark.parametrize("beta", _BETAS) +@pytest.mark.parametrize("pin", [False, True], ids=["default", "portable"]) +def test_aniso_reduces_to_isotropic(beta, pin): + """``U = b/(8*pi^2) I`` must reproduce the isotropic splat, on the live dispatch. + + An analytic identity rather than a comparison against another implementation: a + spherical anisotropic tensor *is* an isotropic B factor, so the two code paths must + agree exactly whatever the cell. It constrains the ``8*pi^2`` conversion and the + diagonal handling of the Mahalanobis form in one assertion, and it fails for a whole + class of index and factor errors that cross-backend parity cannot see because both + sides would share them. + + Ported from ``test_aniso_map_building.py::test_aniso_reduces_to_isotropic``, which + asserted the same identity against ``vectorized_add_to_map_aniso`` -- a kernel no + longer on any dispatch path. + """ + frac, inv_frac, dims, f64 = _cell(beta, dtype=torch.float64) + xyz, adp, occ, A, B = _iso_atoms(f64, n=24, dtype=torch.float64) + voxel = _voxel_size(f64, dims) + grid = torch.zeros(*dims, 3, dtype=torch.float64) + + u_sph = torch.zeros(xyz.shape[0], 6, dtype=torch.float64) + u_sph[:, :3] = (adp / (8.0 * math.pi**2)).unsqueeze(1) + + with (use_portable() if pin else contextlib.nullcontext()): + iso_map = build_electron_density( + grid, xyz, adp, occ, A, B, inv_frac, frac, voxel, dtype=torch.float64 + ) + aniso_map = build_electron_density( + grid, + xyz[:0], adp[:0], occ[:0], A[:0], B[:0], + inv_frac, frac, voxel, + xyz_aniso=xyz, u_aniso=u_sph, occ_aniso=occ, A_aniso=A, B_aniso=B, + dtype=torch.float64, + ) + + rel = _rel_l2(aniso_map, iso_map) + assert rel < 1e-12, ( + f"beta={beta}, {'portable' if pin else 'default'}: a spherical U does not reproduce the iso splat " + f"(rel L2 {rel:.3e}). Suspect the 8*pi^2 conversion or the diagonal of the " + f"Mahalanobis form." + ) + assert float(iso_map.abs().sum()) > 0, "both maps are empty; the identity is vacuous" + + +@pytest.mark.parametrize("n_threads", [1, 2, 4]) +def test_fused_kernel_is_thread_invariant(n_threads): + """The fused splat must give the same map at any thread count. + + It partitions the *output* across threads via ``at::parallel_for`` rather than using + atomics, so a race or a partition-boundary error would show up as a thread-count + dependence. Deliberately non-orthogonal and dense enough that many atoms' spheres + overlap, since a partitioning bug is invisible when no two atoms touch the same voxel. + + **Weak on macOS.** ``_cpp_build.py`` omits ``-fopenmp`` there, and the extension + measurably does not parallelize on this host -- 3000 atoms on a 120x108x96 grid take + 132/131/131 ms at 1/2/4 threads. So this passes locally largely because there is only + one thread to disagree with. It still earns its place: it is real coverage wherever the + extension is built with OpenMP (Linux, CI), and it costs milliseconds. + + Bit-exactness is the right assertion and is not merely aspirational -- verified to hold + on this scene and on a 6x denser one at 2, 4 and 8 threads. Output partitioning means + each voxel is accumulated by one thread over a fixed atom order, so a nonzero + difference is an ordering change worth investigating, not float noise. (When checking + this by hand, compare the maps, not ``dm.sum()``: ``Tensor.sum`` uses a + thread-count-dependent tree reduction and will show a spurious difference of its own.) + + Ported from the ``n_threads``-parametrized thread-safety test in + ``test_cpu_scatter.py``, which exercised the C++ structured scatter -- no longer + reachable from the dispatch. + """ + if not sphere_splat.sphere_splat_available(): + pytest.skip(f"fused CPU sphere splat unavailable: {sphere_splat.last_error()}") + + frac, inv_frac, dims, f64 = _cell(115.0, dtype=torch.float32) + xyz, adp, occ, A, B = _iso_atoms(f64, n=96, dtype=torch.float32, seed=7) + voxel = _voxel_size(f64, dims) + grid = torch.zeros(*dims, 3, dtype=torch.float32) + + original = torch.get_num_threads() + try: + torch.set_num_threads(1) + with contextlib.nullcontext(): + ref = build_electron_density( + grid, xyz, adp, occ, A, B, inv_frac, frac, voxel, dtype=torch.float32 + ) + torch.set_num_threads(n_threads) + with contextlib.nullcontext(): + got = build_electron_density( + grid, xyz, adp, occ, A, B, inv_frac, frac, voxel, dtype=torch.float32 + ) + finally: + torch.set_num_threads(original) + + # Bit-exact: output partitioning means each voxel is summed by exactly one thread in + # the same order regardless of thread count. Any drift is a real ordering change. + assert torch.equal(got, ref), ( + f"{n_threads} threads changed the map (max abs diff " + f"{float((got - ref).abs().max()):.3e}); the fused splat should partition the " + "output, so results must not depend on thread count" + ) + + +def test_fused_gate_requires_one_shared_dtype(): + """The fused kernel takes a *uniform* dtype, not any mix of f32 and f64. + + The C++ selects one ``scalar_t`` from the output map via + ``AT_DISPATCH_FLOATING_TYPES(out.scalar_type(), ...)`` and then reads every other + tensor through a raw pointer of that type. So a float64 map beside float32 atoms would + reinterpret the coordinate buffer as doubles -- garbage values and a 2x out-of-bounds + read -- and the gate has to refuse it. + + Written down because the rule is easy to get wrong when it is restated as a set of + permitted dtypes: "each tensor's dtype is in {f32, f64}" *admits* the mixed case, while + the actual requirement is "all tensors share one dtype drawn from {f32, f64}". The two + read almost identically and only one is memory-safe. In the table that difference is the + ``require_uniform_dtype`` flag, and this asserts it against the row that ships. + """ + from torchref.base.electron_density._backends import DENSITY_BACKENDS + + row = DENSITY_BACKENDS.by_name("cpu_sphere") + f64 = torch.zeros(4, dtype=torch.float64) + f32 = torch.zeros(4, dtype=torch.float32) + six = lambda *ts: list(ts) + [ts[-1]] * (6 - len(ts)) # noqa: E731 + + # Uniform sets satisfy the device/dtype contract. + assert row.mismatch(six(f32, f32, f32)) is None + assert row.mismatch(six(f64, f64, f64)) is None + + # Mixed sets never do, and the message says why. + for mixed in (six(f64, f32, f32), six(f32, f64, f32)): + why = row.mismatch(mixed) + assert why is not None and "single dtype" in why, why + + +def test_fused_extension_compiles(): + """The fused sphere splat must actually build. Fails rather than skipping. + + Every other test in this file -- and in ``tests/unit/structure_factor`` -- calls + ``pytest.skip`` when ``sphere_splat_available()`` is False, which is right for them: + they are testing numerics, and without the extension there is nothing to test. But if + *every* test skips, a build that has stopped working produces an all-green run while + the CPU production path has silently degraded to the portable splat. Dispatch is designed + to degrade quietly, which is correct for users and dangerous for CI. + + So exactly one test asserts the extension builds, and reports the captured diagnostic + plus environment when it does not -- the same stance, and most of the same diagnostic + surface, as the ``TestCompilation`` class in the now-deleted ``test_cpu_scatter.py``. + + Now partly redundant with + ``tests/unit/utils/test_backend_tables.py::test_backend_is_available_where_it_is_expected``, + which generalizes this to every backend from the ``expect_available`` column. This one is + kept for its diagnostics: it prints ninja/CXX/PATH/``TORCH_EXTENSIONS_DIR``, which is what + you actually need when a build breaks. + That guard previously protected the C++ structured scatter, a helper; it now protects + the production CPU splat, so it matters more than it did. + """ + if sphere_splat.sphere_splat_available(): + return + + import os + import shutil + import sys + + err = sphere_splat.last_error() + env_info = ( + f" python: {sys.executable}\n" + f" ninja: {shutil.which('ninja')}\n" + f" CXX env: {os.environ.get('CXX', '')}\n" + f" CC env: {os.environ.get('CC', '')}\n" + f" PATH head: {os.environ.get('PATH', '').split(':')[:5]}\n" + f" TORCH_EXTENSIONS_DIR: " + f"{os.environ.get('TORCH_EXTENSIONS_DIR', '')}\n" + ) + pytest.fail( + "The fused CPU sphere-splat extension failed to build, so CPU dispatch is " + "silently falling back to the portable eager splat for every density " + "calculation.\n" + f"Error: {err}\n\n" + f"Environment:\n{env_info}" + ) + diff --git a/tests/unit/base/test_cpu_scatter.py b/tests/unit/base/test_cpu_scatter.py deleted file mode 100644 index 3341f159..00000000 --- a/tests/unit/base/test_cpu_scatter.py +++ /dev/null @@ -1,577 +0,0 @@ -"""Tests for C++ parallel scatter_add (cpu_scatter.py). - -Verifies correctness and gradient accuracy of structured_scatter_add -against PyTorch's native scatter_add_ for all valid use cases. -""" - -import math - -import pytest -import torch - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - -def _pytorch_scatter_reference(density_cube, wa, wbwc, map_size): - """Reference implementation using PyTorch scatter_add_.""" - C, nx, ny, nz = density_cube.shape - idx_flat = wa[:, :, None, None] + wbwc[:, None, :, :] # (C, nx, ny, nz) - result = torch.zeros(map_size, dtype=density_cube.dtype) - result.scatter_add_(0, idx_flat.reshape(-1), density_cube.reshape(-1)) - return result - - -def _make_valid_indices(C, nx, ny, nz, grid_shape, seed=42): - """Generate valid structured indices that stay within [0, map_size). - - Mimics how production code computes wa and wbwc from center_idx - with modular arithmetic for periodic boundary conditions. - """ - torch.manual_seed(seed) - gx, gy, gz = grid_shape - ny_nz = gy * gz - - # Random center voxels within the grid - center_x = torch.randint(0, gx, (C,)) - center_y = torch.randint(0, gy, (C,)) - center_z = torch.randint(0, gz, (C,)) - - # Axis offsets (symmetric around center, like real code) - half = nx // 2 - offsets = torch.arange(-half, -half + nx) - - # wa: (C, nx) — x-axis indices with PBC wrapping - wa = ((center_x[:, None] + offsets[None, :]) % gx) * ny_nz - - # wbwc: (C, ny, nz) — yz-plane indices with PBC wrapping - half_y = ny // 2 - half_z = nz // 2 - offsets_y = torch.arange(-half_y, -half_y + ny) - offsets_z = torch.arange(-half_z, -half_z + nz) - wb = ((center_y[:, None] + offsets_y[None, :]) % gy) * gz # (C, ny) - wc = ((center_z[:, None] + offsets_z[None, :]) % gz) # (C, nz) - wbwc = wb[:, :, None] + wc[:, None, :] # (C, ny, nz) - - return wa.long(), wbwc.long() - - -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - -@pytest.fixture(scope="module") -def cpp_scatter(): - """Import structured_scatter_add (compiles C++ on first call).""" - try: - from torchref.base.electron_density.kernels.cpu.scatter import structured_scatter_add - return structured_scatter_add - except Exception as e: - pytest.skip(f"C++ scatter not available: {e}") - - -# --------------------------------------------------------------------------- -# Compile-availability test -# --------------------------------------------------------------------------- - -class TestCompilation: - """Verify the C++ extension can be compiled in the current environment. - - This test does NOT skip on failure — failure here is the actual signal - we want from CI / SLURM jobs. The full diagnostic captured by - ``_get_module()`` is included in the assertion message. - """ - - def test_module_compiles(self): - """C++ cpu_scatter extension must build successfully.""" - from torchref.base.electron_density.kernels.cpu import scatter as cpu_scatter - - # Reset any cached failure from a previous test in the same process so - # we get a fresh attempt with up-to-date diagnostics. - cpu_scatter._module_failed = False - cpu_scatter._module_error = None - - mod = cpu_scatter._get_module() - - if mod is None: - err_summary, err_tb = cpu_scatter._module_error or ("unknown", "") - # Surface environment context that commonly differs on SLURM nodes. - import os - import shutil - import sys - env_info = ( - f" python: {sys.executable}\n" - f" ninja: {shutil.which('ninja')}\n" - f" CXX env: {os.environ.get('CXX', '')}\n" - f" CC env: {os.environ.get('CC', '')}\n" - f" PATH head: {os.environ.get('PATH', '').split(':')[:5]}\n" - f" HOME: {os.environ.get('HOME', '')}\n" - f" TORCH_EXTENSIONS_DIR: " - f"{os.environ.get('TORCH_EXTENSIONS_DIR', '')}\n" - ) - pytest.fail( - "C++ cpu_scatter compilation failed.\n" - f"Error: {err_summary}\n\n" - f"Environment:\n{env_info}\n" - f"Full traceback:\n{err_tb}" - ) - - # Sanity-check the built module exposes both entry points - # (one binding per index dtype: int32 fast path + int64 fallback). - for name in ( - "structured_scatter_add_i32", - "structured_scatter_add_i64", - "structured_gather_i32", - "structured_gather_i64", - ): - assert hasattr(mod, name), f"compiled module is missing {name}" - - -# --------------------------------------------------------------------------- -# Correctness tests -# --------------------------------------------------------------------------- - -class TestCorrectnessVsPyTorch: - """Verify C++ scatter matches PyTorch scatter_add_ for all valid cases.""" - - def test_basic_small(self, cpp_scatter): - """Small problem: 4 atoms, 5x5x5 cube, 20^3 grid.""" - C, nx, ny, nz = 4, 5, 5, 5 - grid_shape = (20, 20, 20) - map_size = 20 * 20 * 20 - - torch.manual_seed(0) - density_cube = torch.randn(C, nx, ny, nz) - wa, wbwc = _make_valid_indices(C, nx, ny, nz, grid_shape, seed=0) - - ref = _pytorch_scatter_reference(density_cube, wa, wbwc, map_size) - got = cpp_scatter(density_cube, wa, wbwc, map_size) - - assert got.shape == (map_size,) - assert torch.allclose(ref, got, atol=1e-6), \ - f"max err = {(ref - got).abs().max().item():.2e}" - - def test_single_atom(self, cpp_scatter): - """Edge case: single atom (C=1).""" - C, nx, ny, nz = 1, 7, 7, 7 - grid_shape = (30, 30, 30) - map_size = 30 * 30 * 30 - - torch.manual_seed(1) - density_cube = torch.randn(C, nx, ny, nz) - wa, wbwc = _make_valid_indices(C, nx, ny, nz, grid_shape, seed=1) - - ref = _pytorch_scatter_reference(density_cube, wa, wbwc, map_size) - got = cpp_scatter(density_cube, wa, wbwc, map_size) - - assert torch.allclose(ref, got, atol=1e-6) - - def test_large_chunk(self, cpp_scatter): - """Typical production chunk: 1024 atoms, 17^3 cube.""" - C, nx, ny, nz = 1024, 17, 17, 17 - grid_shape = (216, 90, 72) - map_size = 216 * 90 * 72 - - torch.manual_seed(2) - density_cube = torch.randn(C, nx, ny, nz) - wa, wbwc = _make_valid_indices(C, nx, ny, nz, grid_shape, seed=2) - - ref = _pytorch_scatter_reference(density_cube, wa, wbwc, map_size) - got = cpp_scatter(density_cube, wa, wbwc, map_size) - - assert torch.allclose(ref, got, atol=1e-5), \ - f"max err = {(ref - got).abs().max().item():.2e}" - - def test_asymmetric_cube(self, cpp_scatter): - """Non-cubic ROI: nx != ny != nz.""" - C, nx, ny, nz = 50, 9, 13, 7 - grid_shape = (64, 64, 64) - map_size = 64 ** 3 - - torch.manual_seed(3) - density_cube = torch.randn(C, nx, ny, nz) - wa, wbwc = _make_valid_indices(C, nx, ny, nz, grid_shape, seed=3) - - ref = _pytorch_scatter_reference(density_cube, wa, wbwc, map_size) - got = cpp_scatter(density_cube, wa, wbwc, map_size) - - assert torch.allclose(ref, got, atol=1e-6) - - def test_asymmetric_grid(self, cpp_scatter): - """Non-cubic grid (common in crystallography: a >> b,c).""" - C, nx, ny, nz = 100, 11, 11, 11 - grid_shape = (300, 60, 48) - map_size = 300 * 60 * 48 - - torch.manual_seed(4) - density_cube = torch.randn(C, nx, ny, nz) - wa, wbwc = _make_valid_indices(C, nx, ny, nz, grid_shape, seed=4) - - ref = _pytorch_scatter_reference(density_cube, wa, wbwc, map_size) - got = cpp_scatter(density_cube, wa, wbwc, map_size) - - assert torch.allclose(ref, got, atol=1e-5) - - def test_boundary_wrapping(self, cpp_scatter): - """Atoms near grid boundaries where indices wrap around via PBC.""" - C, nx, ny, nz = 20, 17, 17, 17 - grid_shape = (32, 32, 32) - map_size = 32 ** 3 - - torch.manual_seed(5) - density_cube = torch.randn(C, nx, ny, nz) - - # Force atoms near corners of the grid - gx, gy, gz = grid_shape - ny_nz = gy * gz - half = nx // 2 - offsets = torch.arange(-half, -half + nx) - offsets_y = torch.arange(-half, -half + ny) - offsets_z = torch.arange(-half, -half + nz) - - # Place atoms at grid boundaries (0, last, near-last) - corners = torch.tensor([0, 1, gx - 1, gx - 2, gx // 2]) - center_x = corners.repeat(C // len(corners) + 1)[:C] - center_y = corners.repeat(C // len(corners) + 1)[:C] - center_z = corners.repeat(C // len(corners) + 1)[:C] - - wa = ((center_x[:, None] + offsets[None, :]) % gx) * ny_nz - wb = ((center_y[:, None] + offsets_y[None, :]) % gy) * gz - wc = ((center_z[:, None] + offsets_z[None, :]) % gz) - wbwc = wb[:, :, None] + wc[:, None, :] - - ref = _pytorch_scatter_reference(density_cube, wa, wbwc, map_size) - got = cpp_scatter(density_cube, wa, wbwc, map_size) - - assert torch.allclose(ref, got, atol=1e-6), \ - f"Boundary wrapping failed: max err = {(ref - got).abs().max().item():.2e}" - - def test_overlapping_rois(self, cpp_scatter): - """Multiple atoms with heavily overlapping scatter regions.""" - C, nx, ny, nz = 50, 17, 17, 17 - grid_shape = (40, 40, 40) - map_size = 40 ** 3 - - torch.manual_seed(6) - density_cube = torch.randn(C, nx, ny, nz) - - # All atoms at the same center → maximal overlap - gx, gy, gz = grid_shape - ny_nz = gy * gz - half = nx // 2 - offsets = torch.arange(-half, -half + nx) - offsets_y = torch.arange(-half, -half + ny) - offsets_z = torch.arange(-half, -half + nz) - - center_x = torch.full((C,), gx // 2, dtype=torch.long) - center_y = torch.full((C,), gy // 2, dtype=torch.long) - center_z = torch.full((C,), gz // 2, dtype=torch.long) - - wa = ((center_x[:, None] + offsets[None, :]) % gx) * ny_nz - wb = ((center_y[:, None] + offsets_y[None, :]) % gy) * gz - wc = ((center_z[:, None] + offsets_z[None, :]) % gz) - wbwc = wb[:, :, None] + wc[:, None, :] - - ref = _pytorch_scatter_reference(density_cube, wa, wbwc, map_size) - got = cpp_scatter(density_cube, wa, wbwc, map_size) - - assert torch.allclose(ref, got, atol=1e-5), \ - f"Overlapping ROIs failed: max err = {(ref - got).abs().max().item():.2e}" - - def test_all_positive_values(self, cpp_scatter): - """All density values positive (physical electron density).""" - C, nx, ny, nz = 100, 11, 11, 11 - grid_shape = (64, 64, 64) - map_size = 64 ** 3 - - torch.manual_seed(7) - density_cube = torch.rand(C, nx, ny, nz) * 10.0 # strictly positive - wa, wbwc = _make_valid_indices(C, nx, ny, nz, grid_shape, seed=7) - - ref = _pytorch_scatter_reference(density_cube, wa, wbwc, map_size) - got = cpp_scatter(density_cube, wa, wbwc, map_size) - - assert torch.allclose(ref, got, atol=1e-5) - # Result should also be non-negative - assert got.min() >= -1e-6 - - def test_zero_density(self, cpp_scatter): - """All-zero density cube (should produce zero map).""" - C, nx, ny, nz = 10, 5, 5, 5 - grid_shape = (20, 20, 20) - map_size = 20 ** 3 - - density_cube = torch.zeros(C, nx, ny, nz) - wa, wbwc = _make_valid_indices(C, nx, ny, nz, grid_shape, seed=8) - - got = cpp_scatter(density_cube, wa, wbwc, map_size) - assert torch.all(got == 0.0) - - def test_sparse_density(self, cpp_scatter): - """Mostly-zero density with a few nonzero entries.""" - C, nx, ny, nz = 32, 9, 9, 9 - grid_shape = (50, 50, 50) - map_size = 50 ** 3 - - torch.manual_seed(9) - density_cube = torch.zeros(C, nx, ny, nz) - # Set only ~1% of entries to nonzero - mask = torch.rand_like(density_cube) < 0.01 - density_cube[mask] = torch.randn(mask.sum()) - - wa, wbwc = _make_valid_indices(C, nx, ny, nz, grid_shape, seed=9) - - ref = _pytorch_scatter_reference(density_cube, wa, wbwc, map_size) - got = cpp_scatter(density_cube, wa, wbwc, map_size) - - assert torch.allclose(ref, got, atol=1e-6) - - def test_large_values(self, cpp_scatter): - """Large magnitude values (stress test float32 accumulation).""" - C, nx, ny, nz = 200, 11, 11, 11 - grid_shape = (80, 80, 80) - map_size = 80 ** 3 - - torch.manual_seed(10) - density_cube = torch.randn(C, nx, ny, nz) * 1e4 - wa, wbwc = _make_valid_indices(C, nx, ny, nz, grid_shape, seed=10) - - ref = _pytorch_scatter_reference(density_cube, wa, wbwc, map_size) - got = cpp_scatter(density_cube, wa, wbwc, map_size) - - # Looser tolerance for large values - assert torch.allclose(ref, got, rtol=1e-5, atol=1e-2), \ - f"Large values: max err = {(ref - got).abs().max().item():.2e}" - - def test_sorted_atoms(self, cpp_scatter): - """Atoms sorted by 1D center (production configuration).""" - C, nx, ny, nz = 512, 17, 17, 17 - grid_shape = (216, 90, 72) - map_size = 216 * 90 * 72 - - torch.manual_seed(11) - density_cube = torch.randn(C, nx, ny, nz) - wa, wbwc = _make_valid_indices(C, nx, ny, nz, grid_shape, seed=11) - - # Sort by first wa entry (approximates sorting by 1D center) - order = torch.argsort(wa[:, wa.shape[1] // 2]) - density_cube = density_cube[order] - wa = wa[order] - wbwc = wbwc[order] - - ref = _pytorch_scatter_reference(density_cube, wa, wbwc, map_size) - got = cpp_scatter(density_cube, wa, wbwc, map_size) - - assert torch.allclose(ref, got, atol=1e-5) - - def test_remainder_chunk(self, cpp_scatter): - """Non-power-of-2 atom count (tests remainder handling).""" - C, nx, ny, nz = 37, 9, 9, 9 - grid_shape = (50, 50, 50) - map_size = 50 ** 3 - - torch.manual_seed(12) - density_cube = torch.randn(C, nx, ny, nz) - wa, wbwc = _make_valid_indices(C, nx, ny, nz, grid_shape, seed=12) - - ref = _pytorch_scatter_reference(density_cube, wa, wbwc, map_size) - got = cpp_scatter(density_cube, wa, wbwc, map_size) - - assert torch.allclose(ref, got, atol=1e-6) - - -# --------------------------------------------------------------------------- -# Gradient tests -# --------------------------------------------------------------------------- - -class TestGradients: - """Verify gradient correctness of structured_scatter_add.""" - - def test_gradcheck_small(self, cpp_scatter): - """Numerical gradient check via finite differences (float32). - - Uses a small problem and generous tolerances appropriate for float32. - """ - C, nx, ny, nz = 3, 5, 5, 5 - grid_shape = (20, 20, 20) - map_size = 20 ** 3 - - torch.manual_seed(20) - wa, wbwc = _make_valid_indices(C, nx, ny, nz, grid_shape, seed=20) - - # Compute analytical gradient - dc = torch.randn(C, nx, ny, nz, requires_grad=True) - grad_output = torch.randn(map_size) - result = cpp_scatter(dc, wa, wbwc, map_size) - result.backward(grad_output) - analytical = dc.grad.clone() - - # Compute numerical gradient via finite differences - eps = 1e-4 - numerical = torch.zeros_like(dc) - dc_flat = dc.data.view(-1) - for i in range(dc_flat.numel()): - dc_flat[i] += eps - f_plus = cpp_scatter(dc.data.view(C, nx, ny, nz), wa, wbwc, map_size) - dc_flat[i] -= 2 * eps - f_minus = cpp_scatter(dc.data.view(C, nx, ny, nz), wa, wbwc, map_size) - dc_flat[i] += eps # restore - numerical.view(-1)[i] = ((f_plus - f_minus) * grad_output).sum() / (2 * eps) - - assert torch.allclose(analytical, numerical, atol=1e-3, rtol=1e-3), \ - f"Gradcheck failed: max err = {(analytical - numerical).abs().max().item():.2e}" - - def test_gradient_vs_pytorch(self, cpp_scatter): - """Compare gradients from C++ scatter vs PyTorch scatter.""" - C, nx, ny, nz = 64, 11, 11, 11 - grid_shape = (64, 64, 64) - map_size = 64 ** 3 - - torch.manual_seed(21) - wa, wbwc = _make_valid_indices(C, nx, ny, nz, grid_shape, seed=21) - - # Random upstream gradient - grad_output = torch.randn(map_size) - - # PyTorch path - dc_pt = torch.randn(C, nx, ny, nz, requires_grad=True) - ref = _pytorch_scatter_reference(dc_pt, wa, wbwc, map_size) - ref.backward(grad_output) - grad_pt = dc_pt.grad.clone() - - # C++ path - dc_cpp = dc_pt.data.clone().requires_grad_(True) - got = cpp_scatter(dc_cpp, wa, wbwc, map_size) - got.backward(grad_output) - grad_cpp = dc_cpp.grad.clone() - - assert torch.allclose(grad_pt, grad_cpp, atol=1e-6), \ - f"Gradient mismatch: max err = {(grad_pt - grad_cpp).abs().max().item():.2e}" - - def test_gradient_accumulation(self, cpp_scatter): - """Gradients accumulate correctly across multiple scatter calls (like chunk loop).""" - C, nx, ny, nz = 32, 9, 9, 9 - grid_shape = (40, 40, 40) - map_size = 40 ** 3 - - torch.manual_seed(22) - wa, wbwc = _make_valid_indices(C, nx, ny, nz, grid_shape, seed=22) - - # Two chunks summed (simulates the chunk loop) - dc1 = torch.randn(C, nx, ny, nz, requires_grad=True) - dc2 = torch.randn(C, nx, ny, nz, requires_grad=True) - - result = cpp_scatter(dc1, wa, wbwc, map_size) + \ - cpp_scatter(dc2, wa, wbwc, map_size) - loss = result.sum() - loss.backward() - - # Reference: expected gradient of sum() through gather is all-ones gathered - # grad_dc[c, ix, iy, iz] = 1.0 for all entries (since loss = sum of all) - expected_grad = torch.ones(C, nx, ny, nz) - - assert torch.allclose(dc1.grad, expected_grad, atol=1e-6), \ - f"dc1 grad err: {(dc1.grad - expected_grad).abs().max().item():.2e}" - assert torch.allclose(dc2.grad, expected_grad, atol=1e-6), \ - f"dc2 grad err: {(dc2.grad - expected_grad).abs().max().item():.2e}" - - def test_gradient_with_downstream_loss(self, cpp_scatter): - """Gradient through a realistic loss: MSE between scattered result and target.""" - C, nx, ny, nz = 16, 7, 7, 7 - grid_shape = (32, 32, 32) - map_size = 32 ** 3 - - torch.manual_seed(23) - wa, wbwc = _make_valid_indices(C, nx, ny, nz, grid_shape, seed=23) - target = torch.randn(map_size) - - # C++ path - dc_cpp = torch.randn(C, nx, ny, nz, requires_grad=True) - result_cpp = cpp_scatter(dc_cpp, wa, wbwc, map_size) - loss_cpp = ((result_cpp - target) ** 2).mean() - loss_cpp.backward() - grad_cpp = dc_cpp.grad.clone() - - # PyTorch reference - dc_pt = dc_cpp.data.clone().requires_grad_(True) - idx_flat = wa[:, :, None, None] + wbwc[:, None, :, :] - result_pt = torch.zeros(map_size) - result_pt.scatter_add_(0, idx_flat.reshape(-1), dc_pt.reshape(-1)) - loss_pt = ((result_pt - target) ** 2).mean() - loss_pt.backward() - grad_pt = dc_pt.grad.clone() - - assert abs(loss_cpp.item() - loss_pt.item()) < 1e-5, \ - f"Loss mismatch: {loss_cpp.item():.6f} vs {loss_pt.item():.6f}" - assert torch.allclose(grad_cpp, grad_pt, atol=1e-5), \ - f"Gradient mismatch: max err = {(grad_cpp - grad_pt).abs().max().item():.2e}" - - def test_gradient_overlapping_indices(self, cpp_scatter): - """Gradient correct when multiple atoms scatter to the same voxel.""" - C, nx, ny, nz = 20, 11, 11, 11 - grid_shape = (30, 30, 30) - map_size = 30 ** 3 - - torch.manual_seed(24) - - # All atoms at the same center → all scatter to same voxels - gx, gy, gz = grid_shape - ny_nz = gy * gz - half = nx // 2 - offsets = torch.arange(-half, -half + nx) - offsets_y = torch.arange(-half, -half + ny) - offsets_z = torch.arange(-half, -half + nz) - - center = torch.full((C,), gx // 2, dtype=torch.long) - wa = ((center[:, None] + offsets[None, :]) % gx) * ny_nz - wb = ((center[:, None] + offsets_y[None, :]) % gy) * gz - wc = ((center[:, None] + offsets_z[None, :]) % gz) - wbwc = wb[:, :, None] + wc[:, None, :] - - grad_output = torch.randn(map_size) - - # C++ path - dc_cpp = torch.randn(C, nx, ny, nz, requires_grad=True) - result_cpp = cpp_scatter(dc_cpp, wa, wbwc, map_size) - result_cpp.backward(grad_output) - grad_cpp = dc_cpp.grad.clone() - - # PyTorch reference - dc_pt = dc_cpp.data.clone().requires_grad_(True) - ref = _pytorch_scatter_reference(dc_pt, wa, wbwc, map_size) - ref.backward(grad_output) - grad_pt = dc_pt.grad.clone() - - assert torch.allclose(grad_cpp, grad_pt, atol=1e-6), \ - f"Overlapping gradient err: {(grad_cpp - grad_pt).abs().max().item():.2e}" - - -# --------------------------------------------------------------------------- -# Thread safety -# --------------------------------------------------------------------------- - -class TestThreadSafety: - """Verify results are independent of thread count.""" - - @pytest.mark.parametrize("n_threads", [1, 2, 4, 8]) - def test_thread_count_invariance(self, cpp_scatter, n_threads): - """Result must be identical regardless of OMP thread count.""" - C, nx, ny, nz = 200, 13, 13, 13 - grid_shape = (100, 80, 60) - map_size = 100 * 80 * 60 - - torch.manual_seed(30) - density_cube = torch.randn(C, nx, ny, nz) - wa, wbwc = _make_valid_indices(C, nx, ny, nz, grid_shape, seed=30) - - ref = _pytorch_scatter_reference(density_cube, wa, wbwc, map_size) - - prev_threads = torch.get_num_threads() - try: - torch.set_num_threads(n_threads) - got = cpp_scatter(density_cube, wa, wbwc, map_size) - assert torch.allclose(ref, got, atol=1e-5), \ - f"{n_threads} threads: max err = {(ref - got).abs().max().item():.2e}" - finally: - torch.set_num_threads(prev_threads) diff --git a/tests/unit/base/test_ds_dispatch.py b/tests/unit/base/test_ds_dispatch.py deleted file mode 100644 index 2416c4e3..00000000 --- a/tests/unit/base/test_ds_dispatch.py +++ /dev/null @@ -1,127 +0,0 @@ -"""Unit tests for the direct-summation dispatch / checkpointed backend. - -CPU-only (no GPU required). Verifies that the recompute-on-backward -checkpointed backend matches the eager-autograd reference to machine -precision, that reflection-chunking is exact, that gradients pass -``gradcheck``, and that the engine guards behave. -""" - -import pytest -import torch - -from torchref.base.direct_summation import Engine, ds_aniso, ds_iso -from torchref.base.direct_summation import dispatch as D -from torchref.base.direct_summation.anisotropic import aniso_structure_factor_torched -from torchref.base.direct_summation.isotropic import iso_structure_factor_torched -from torchref.config import dtypes - -pytestmark = pytest.mark.unit - -# The eager reference functions internally cast hkl to ``dtypes.float`` so the -# eager-parity comparisons run in the configured float dtype (float32 by -# default). gradcheck, which never calls the eager fn, runs in float64. -_F = dtypes.float -_eager_atol = 1e-4 if _F == torch.float32 else 1e-10 - - -def _p1(c): - return c.unsqueeze(2) - - -def _inputs(N=4, R=7, seed=0, dtype=None): - dtype = dtype or _F - torch.manual_seed(seed) - hkl = torch.randint(-3, 4, (R, 3)).to(dtype) - s = torch.rand(R, dtype=dtype) * 0.5 - svec = torch.randn(R, 3, dtype=dtype) * 0.3 - A = torch.rand(N, 5, dtype=dtype) - B = torch.rand(N, 5, dtype=dtype) + 0.5 - return hkl, s, svec, A, B - - -def _leaves(N=4, seed=1, dtype=None): - dtype = dtype or _F - torch.manual_seed(seed) - xyz = torch.rand(N, 3, dtype=dtype, requires_grad=True) - occ = (torch.rand(N, dtype=dtype) * 0.4 + 0.6).requires_grad_() - adp = (torch.rand(N, dtype=dtype) * 10 + 5).requires_grad_() - U = (torch.rand(N, 6, dtype=dtype) * 0.04 + 0.01).requires_grad_() - return xyz, occ, adp, U - - -def test_checkpointed_iso_matches_eager(): - hkl, s, _, A, B = _inputs() - xyz, occ, adp, _ = _leaves() - F = D._checkpointed_iso(hkl, s, xyz, occ, adp, A, B, max_memory_gb=None) - ((F.real**2 + 2 * F.imag).sum()).backward() - gck = (xyz.grad.clone(), occ.grad.clone(), adp.grad.clone()) - - xyz2, occ2, adp2, _ = _leaves() - Fe = iso_structure_factor_torched( - hkl=hkl, s=s, xyz_fractional=xyz2, occ=occ2, scattering_factors=None, - adp=adp2, spacegroup=_p1, max_memory_gb=2.0, A=A, B_coeff=B, - ) - ((Fe.real**2 + 2 * Fe.imag).sum()).backward() - assert torch.allclose(F, Fe, atol=_eager_atol, rtol=1e-3) - for g, ge in zip(gck, (xyz2.grad, occ2.grad, adp2.grad)): - assert torch.allclose(g, ge, atol=_eager_atol, rtol=1e-3) - - -def test_checkpointed_aniso_matches_eager(): - hkl, _, svec, A, B = _inputs() - xyz, occ, _, U = _leaves() - F = D._checkpointed_aniso(hkl, svec, xyz, occ, U, A, B, max_memory_gb=None) - ((F.real**2 + 2 * F.imag).sum()).backward() - gck = (xyz.grad.clone(), occ.grad.clone(), U.grad.clone()) - - xyz2, occ2, _, U2 = _leaves() - Fe = aniso_structure_factor_torched( - hkl=hkl, s_vector=svec, xyz_fractional=xyz2, occ=occ2, - scattering_factors=None, U=U2, spacegroup=_p1, max_memory_gb=2.0, - A=A, B_coeff=B, - ) - ((Fe.real**2 + 2 * Fe.imag).sum()).backward() - assert torch.allclose(F, Fe, atol=_eager_atol, rtol=1e-3) - for g, ge in zip(gck, (xyz2.grad, occ2.grad, U2.grad)): - assert torch.allclose(g, ge, atol=_eager_atol, rtol=1e-3) - - -def test_checkpointed_chunking_is_exact(): - # Chunking slices the reflection dimension, so each reflection's atom sum is - # computed within a single chunk -- the result is mathematically identical - # regardless of chunk size. It is not, however, guaranteed bit-exact: the - # phase matmul dispatches to different BLAS kernels for (R,3)x(3,N) vs - # (1,3)x(3,N), and last-ULP rounding there is CPU/BLAS-dependent. Assert - # numerical equivalence, not bitwise equality. - hkl, s, _, A, B = _inputs(R=11, dtype=torch.float64) - xyz, occ, adp, _ = _leaves(dtype=torch.float64) - F_full = D._checkpointed_iso(hkl, s, xyz, occ, adp, A, B, max_memory_gb=None) - F_chunk = D._checkpointed_iso(hkl, s, xyz, occ, adp, A, B, max_memory_gb=1e-7) - assert torch.allclose(F_full, F_chunk, rtol=1e-10, atol=1e-12) - - -def test_checkpointed_iso_gradcheck(): - hkl, s, _, A, B = _inputs(dtype=torch.float64) - xyz, occ, adp, _ = _leaves(dtype=torch.float64) - - def f(x, o, a): - return D._checkpointed_iso(hkl, s, x, o, a, A, B, max_memory_gb=1e-7) - - assert torch.autograd.gradcheck(f, (xyz, occ, adp), eps=1e-6, atol=1e-5) - - -def test_empty_atoms_returns_zeros(): - hkl, s, _, _, _ = _inputs() - empty = torch.zeros(0, 3, dtype=torch.float64) - z = torch.zeros(0, dtype=torch.float64) - z5 = torch.zeros(0, 5, dtype=torch.float64) - F = ds_iso(hkl, s, empty, z, z, z5, z5, engine=Engine.AUTO) - assert F.shape == (hkl.shape[0],) - assert F.abs().sum().item() == 0.0 - - -def test_explicit_triton_engine_rejects_cpu(): - hkl, s, _, A, B = _inputs() - xyz, occ, adp, _ = _leaves() - with pytest.raises(RuntimeError): - ds_iso(hkl, s, xyz, occ, adp, A, B, engine=Engine.TRITON) diff --git a/tests/unit/base/test_kernel_import_shims.py b/tests/unit/base/test_kernel_import_shims.py index 6f1bb2c6..64839817 100644 --- a/tests/unit/base/test_kernel_import_shims.py +++ b/tests/unit/base/test_kernel_import_shims.py @@ -11,6 +11,12 @@ import pytest +from torchref.utils.backends import triton_available + +_needs_triton = pytest.mark.skipif( + not triton_available(), reason="CUDA/Triton backend requires the triton package" +) + def test_kernels_public_api_resolves(): kernels = importlib.import_module("torchref.base.kernels") @@ -53,19 +59,20 @@ def test_math_torch_legacy_reexports_resolve(): def test_main_namespace_preserves_moved_symbols(): - """``main`` re-exports the shared splat helpers so its namespace is stable. + """``main`` keeps the names other modules and tests reach for. + + Only two things need to resolve here now. ``_get_radius_offsets`` because + ``torchref.scaling.solvent`` imports it from ``main`` rather than from its defining + module, and the two dispatchers because they are ``main``'s own API. - These are the LIVE helpers reused by the variable-radius kernels (and - ``_get_radius_offsets`` by ``torchref.scaling.solvent``); the old fixed-radius - entry points were removed in the kernel cleanup. + The list used to also carry ``_do_structured_scatter``, ``_get_cpp_scatter``, + ``_separable_density`` and ``_aniso_density_cube``. Those were re-exported for a + grouped-separable/cube splat chain that the fused sphere kernel superseded; nothing + called it, and it has been deleted along with the four modules behind it. """ main = importlib.import_module("torchref.base.electron_density.main") moved = [ "_get_radius_offsets", - "_do_structured_scatter", - "_get_cpp_scatter", - "_separable_density", - "_aniso_density_cube", # dispatchers stay defined here "_add_isotropic", "_add_anisotropic", @@ -84,14 +91,23 @@ def test_solvent_radius_offsets_import_path(): [ "torchref.base.electron_density.kernels", "torchref.base.electron_density.kernels.offsets", - "torchref.base.electron_density.kernels.cpu.separable", - "torchref.base.electron_density.kernels.cpu.aniso", - "torchref.base.electron_density.kernels.cpu.scatter", - "torchref.base.electron_density.kernels.cpu.scatter_dispatch", "torchref.base.electron_density.kernels.cpu.jit_reference", "torchref.base.electron_density.kernels.cpu.variable_radius", - "torchref.base.electron_density.kernels.cuda.fused", - "torchref.base.electron_density.kernels.cuda.variable_radius", + # CUDA/Triton backend: importing pulls in `triton`, absent on non-CUDA + # hosts (e.g. macOS), so gate these on triton availability. + pytest.param( + "torchref.base.electron_density.kernels.cuda.fused", marks=_needs_triton + ), + pytest.param( + "torchref.base.electron_density.kernels.cuda.variable_radius", + marks=_needs_triton, + ), + # MPS Metal kernels: importing must not trigger shader compilation, so + # these resolve cleanly on every platform (compile is deferred to first use). + "torchref.base.electron_density.kernels.mps", + "torchref.base.electron_density.kernels.mps.compile", + "torchref.base.electron_density.kernels.mps.variable_radius", + "torchref.base.electron_density.kernels.mps._shaders", ], ) def test_new_kernel_modules_import(modname): diff --git a/tests/unit/base/test_math_torch.py b/tests/unit/base/test_math_torch.py index 9ac12c10..3bf19be0 100644 --- a/tests/unit/base/test_math_torch.py +++ b/tests/unit/base/test_math_torch.py @@ -92,8 +92,8 @@ def test_coordinate_transforms_gpu(self, mock_cell, random_coordinates, gpu_devi frac = cartesian_to_fractional_torch(coords, cell) cart_back = fractional_to_cartesian_torch(frac, cell) - assert frac.device.type == "cuda" - assert cart_back.device.type == "cuda" + assert frac.device.type == gpu_device.type + assert cart_back.device.type == gpu_device.type class TestGridFunctions: diff --git a/tests/unit/base/test_max_eig_sym3.py b/tests/unit/base/test_max_eig_sym3.py new file mode 100644 index 00000000..3bc1da5d --- /dev/null +++ b/tests/unit/base/test_max_eig_sym3.py @@ -0,0 +1,72 @@ +"""Closed-form largest eigenvalue of a symmetric 3x3 vs torch.linalg.eigvalsh. + +``_max_eig_sym3`` replaces ``eigvalsh`` in the anisotropic splat-radius policy so +it runs natively on MPS; it must agree with the reference eigendecomposition, +including on diagonal and near-degenerate matrices. +""" + +import pytest +import torch + +from torchref.base.electron_density.radius_policy import ( + _max_eig_sym3, + _u6_to_u3, + per_atom_radius_aniso, +) + +pytestmark = pytest.mark.unit + + +def _sym(n, seed, scale=1.0): + g = torch.Generator().manual_seed(seed) + M = torch.randn(n, 3, 3, generator=g) * scale + return (M + M.transpose(1, 2)) / 2 + + +def test_matches_eigvalsh_random(): + S = _sym(4000, 0) + ref = torch.linalg.eigvalsh(S).max(dim=1).values + got = _max_eig_sym3(S) + assert torch.allclose(got, ref, atol=1e-4, rtol=1e-4) + + +def test_matches_eigvalsh_diagonal(): + # purely diagonal (off-diagonals zero) -> largest diagonal entry + d = torch.tensor([[3.0, -1.0, 2.0], [5.0, 5.0, 5.0], [0.1, 0.2, 0.05]]) + S = torch.diag_embed(d) + ref = torch.linalg.eigvalsh(S).max(dim=1).values + got = _max_eig_sym3(S) + assert torch.allclose(got, ref, atol=1e-5) + assert torch.allclose(got, d.max(dim=1).values, atol=1e-5) + + +def test_matches_eigvalsh_near_degenerate(): + # near-isotropic (all eigenvalues nearly equal) stresses the p2->0 branch + S = torch.eye(3).expand(50, 3, 3).clone() + S += _sym(50, 7, scale=1e-4) + ref = torch.linalg.eigvalsh(S).max(dim=1).values + got = _max_eig_sym3(S) + assert torch.allclose(got, ref, atol=1e-4) + + +def test_radius_matches_eigvalsh_path(): + # The (quantized) aniso radius must be identical to the eigvalsh-based one. + g = torch.Generator().manual_seed(3) + n = 500 + u = torch.zeros(n, 6) + u[:, :3] = torch.rand(n, 3, generator=g) * 0.12 + 0.02 + u[:, 3:] = (torch.rand(n, 3, generator=g) - 0.5) * 0.02 + B = torch.rand(n, 5, generator=g) * 20 + 2 + + rad_closed = per_atom_radius_aniso(B, u, n_sigma=3.0) + + # Reference radius using eigvalsh directly. + from torchref.base.electron_density.radius_policy import ( + EIGHT_PI2, _ceil_round, R_LO, R_HI, + ) + b_form = B.max(dim=1).values + lam = torch.linalg.eigvalsh(_u6_to_u3(u)).max(dim=1).values + sigma = torch.sqrt((b_form / EIGHT_PI2 + lam).clamp(min=1e-6)) + rad_ref = _ceil_round(3.0 * sigma).clamp(min=R_LO, max=R_HI) + + assert torch.equal(rad_closed, rad_ref) diff --git a/tests/unit/base/test_vectorized_add_to_map.py b/tests/unit/base/test_vectorized_add_to_map.py deleted file mode 100644 index 56244512..00000000 --- a/tests/unit/base/test_vectorized_add_to_map.py +++ /dev/null @@ -1,446 +0,0 @@ -"""Tests for vectorized_add_to_map with automatic compilation.""" -import pytest -import torch -import numpy as np -import time - -from torchref.base.kernels import ( - vectorized_add_to_map, - compute_metric_tensor, - precompute_fractional_coords, - warmup, -) -from torchref.base.math_torch import ( - vectorized_add_to_map as original_add_to_map, -) - - -class TestValidityAgainstOriginal: - """Compare compiled version against original math_torch implementation.""" - - @pytest.fixture - def setup_tensors(self): - """Create test tensors matching both function signatures.""" - def _setup(n_atoms=100, n_voxels=500, grid_shape=(64, 64, 64), - device="cpu", dtype=torch.float64, seed=42): - torch.manual_seed(seed) - - # Cell parameters (orthorhombic for simplicity) - cell = torch.tensor( - [50.0, 60.0, 70.0, 90.0, 90.0, 90.0], dtype=dtype, device=device - ) - frac_matrix = torch.diag(cell[:3]) - inv_frac_matrix = torch.diag(1.0 / cell[:3]) - - # Atom parameters - xyz = torch.rand(n_atoms, 3, device=device, dtype=dtype) * 30 + 10 - b = torch.rand(n_atoms, device=device, dtype=dtype) * 40 + 10 - A = torch.rand(n_atoms, 5, device=device, dtype=dtype) * 5 + 1 - B = torch.rand(n_atoms, 5, device=device, dtype=dtype) * 10 + 2 - occ = torch.rand(n_atoms, device=device, dtype=dtype) * 0.5 + 0.5 - - # Voxel coordinates (Cartesian, around atoms) - surrounding_coords = ( - xyz[:, None, :] - + torch.randn(n_atoms, n_voxels, 3, device=device, dtype=dtype) * 2 - ) - voxel_indices = torch.randint( - 0, min(grid_shape), (n_atoms, n_voxels, 3), device=device - ) - - return { - "surrounding_coords": surrounding_coords, - "voxel_indices": voxel_indices, - "grid_shape": grid_shape, - "xyz": xyz, - "b": b, - "frac_matrix": frac_matrix, - "inv_frac_matrix": inv_frac_matrix, - "A": A, - "B": B, - "occ": occ, - } - return _setup - - def test_output_correlation_with_original(self, setup_tensors): - """Compiled output correlates highly with original. - - Note: The compiled version uses softplus for smooth gradients while - the original uses hard clamp. This gives slightly different B_total - values and thus different outputs, but they should be highly correlated. - """ - data = setup_tensors(n_atoms=50, n_voxels=200) - - # Original - map_orig = torch.zeros(data["grid_shape"], dtype=data["xyz"].dtype) - map_orig = original_add_to_map( - data["surrounding_coords"], data["voxel_indices"], map_orig, - data["xyz"], data["b"], data["inv_frac_matrix"], data["frac_matrix"], - data["A"], data["B"], data["occ"] - ) - - # Auto-compiled version (same API now) - map_compiled = torch.zeros(data["grid_shape"], dtype=data["xyz"].dtype) - map_compiled = vectorized_add_to_map( - data["surrounding_coords"], data["voxel_indices"], map_compiled, - data["xyz"], data["b"], data["inv_frac_matrix"], data["frac_matrix"], - data["A"], data["B"], data["occ"] - ) - - # Check correlation instead of exact match (softplus vs clamp causes differences) - orig_flat = map_orig.flatten() - compiled_flat = map_compiled.flatten() - - # Remove zeros for correlation - mask = (orig_flat != 0) | (compiled_flat != 0) - if mask.sum() > 0: - correlation = torch.corrcoef( - torch.stack([orig_flat[mask], compiled_flat[mask]]) - )[0, 1] - assert correlation > 0.999, f"Correlation too low: {correlation}" - - # Check max relative difference is reasonable (within 1%) - max_rel_diff = ( - (map_orig - map_compiled).abs().max() / (map_orig.abs().max() + 1e-10) - ) - assert max_rel_diff < 0.01, f"Max relative difference too high: {max_rel_diff}" - - @pytest.mark.parametrize("n_atoms,n_voxels", [ - (10, 100), - (100, 500), - (500, 1000) - ]) - def test_dynamic_shapes(self, setup_tensors, n_atoms, n_voxels): - """Different sizes should work without errors.""" - data = setup_tensors(n_atoms=n_atoms, n_voxels=n_voxels) - - map_out = torch.zeros(data["grid_shape"], dtype=data["xyz"].dtype) - result = vectorized_add_to_map( - data["surrounding_coords"], data["voxel_indices"], map_out, - data["xyz"], data["b"], data["inv_frac_matrix"], data["frac_matrix"], - data["A"], data["B"], data["occ"] - ) - - assert result.shape == data["grid_shape"] - assert torch.isfinite(result).all() - - def test_output_in_place(self, setup_tensors): - """Verify that the function modifies the map in-place (like the original).""" - data = setup_tensors(n_atoms=20, n_voxels=100) - - map_in = torch.zeros(data["grid_shape"], dtype=data["xyz"].dtype) - - result = vectorized_add_to_map( - data["surrounding_coords"], data["voxel_indices"], map_in, - data["xyz"], data["b"], data["inv_frac_matrix"], data["frac_matrix"], - data["A"], data["B"], data["occ"] - ) - - # Function modifies in-place and returns same tensor (like original scatter_add_nd) - assert result is map_in - # Map should now contain density values - assert result.abs().sum() > 0 - - -class TestGradientValidity: - """Test gradient flow through compiled function.""" - - @pytest.fixture - def grad_tensors(self): - """Create tensors with requires_grad=True.""" - def _setup(n_atoms=50, n_voxels=200, grid_shape=(32, 32, 32), seed=42): - torch.manual_seed(seed) - dtype = torch.float64 - device = "cpu" - - frac_matrix = torch.diag( - torch.tensor([50.0, 60.0, 70.0], dtype=dtype, device=device) - ) - inv_frac_matrix = torch.diag( - 1.0 / torch.tensor([50.0, 60.0, 70.0], dtype=dtype, device=device) - ) - - xyz = torch.rand(n_atoms, 3, dtype=dtype, device=device).requires_grad_(True) - b = (torch.rand(n_atoms, dtype=dtype, device=device) * 40 + 10).requires_grad_(True) - A = torch.rand(n_atoms, 5, dtype=dtype, device=device).requires_grad_(True) - B = (torch.rand(n_atoms, 5, dtype=dtype, device=device) * 10 + 2).requires_grad_(True) - occ = torch.rand(n_atoms, dtype=dtype, device=device).requires_grad_(True) - - surrounding_coords = ( - xyz.detach()[:, None, :] - + torch.randn(n_atoms, n_voxels, 3, dtype=dtype, device=device) * 2 - ) - voxel_indices = torch.randint( - 0, min(grid_shape), (n_atoms, n_voxels, 3), device=device - ) - - return { - "surrounding_coords": surrounding_coords, - "voxel_indices": voxel_indices, - "grid_shape": grid_shape, - "xyz": xyz, - "b": b, - "frac_matrix": frac_matrix, - "inv_frac_matrix": inv_frac_matrix, - "A": A, - "B": B, - "occ": occ, - } - return _setup - - def test_gradients_exist(self, grad_tensors): - """All gradient-requiring params should have gradients after backward.""" - data = grad_tensors() - - density_map = torch.zeros(data["grid_shape"], dtype=data["xyz"].dtype) - result = vectorized_add_to_map( - data["surrounding_coords"], data["voxel_indices"], density_map, - data["xyz"], data["b"], data["inv_frac_matrix"], data["frac_matrix"], - data["A"], data["B"], data["occ"] - ) - - loss = result.sum() - loss.backward() - - # Check all gradients exist and are finite - for name in ["xyz", "b", "A", "B", "occ"]: - grad = data[name].grad - assert grad is not None, f"Gradient for {name} is None" - assert torch.isfinite(grad).all(), ( - f"Gradient for {name} has non-finite values" - ) - assert grad.abs().sum() > 0, f"Gradient for {name} is all zeros" - - def test_gradients_match_original(self, grad_tensors): - """Gradients from compiled version match original implementation.""" - data = grad_tensors(n_atoms=30, n_voxels=100) - - # Clone tensors for original - xyz_orig = data["xyz"].detach().clone().requires_grad_(True) - b_orig = data["b"].detach().clone().requires_grad_(True) - A_orig = data["A"].detach().clone().requires_grad_(True) - B_orig = data["B"].detach().clone().requires_grad_(True) - occ_orig = data["occ"].detach().clone().requires_grad_(True) - - # Create same surrounding coords for both - torch.manual_seed(123) - surrounding_coords = ( - xyz_orig.detach()[:, None, :] - + torch.randn(30, 100, 3, dtype=data["xyz"].dtype) * 2 - ) - - # Original backward - map_orig = torch.zeros(data["grid_shape"], dtype=xyz_orig.dtype) - map_orig = original_add_to_map( - surrounding_coords, data["voxel_indices"], map_orig, - xyz_orig, b_orig, data["inv_frac_matrix"], data["frac_matrix"], - A_orig, B_orig, occ_orig - ) - loss_orig = map_orig.sum() - loss_orig.backward() - - # Compiled backward - map_compiled = torch.zeros(data["grid_shape"], dtype=data["xyz"].dtype) - map_compiled = vectorized_add_to_map( - surrounding_coords, data["voxel_indices"], map_compiled, - data["xyz"], data["b"], data["inv_frac_matrix"], data["frac_matrix"], - data["A"], data["B"], data["occ"] - ) - loss_compiled = map_compiled.sum() - loss_compiled.backward() - - # Compare gradients - looser tolerance due to softplus vs clamp differences - assert torch.allclose( - data["A"].grad, A_orig.grad, rtol=5e-2, atol=1e-3 - ), "A gradients don't match" - assert torch.allclose( - data["B"].grad, B_orig.grad, rtol=5e-2, atol=1e-3 - ), "B gradients don't match" - assert torch.allclose( - data["occ"].grad, occ_orig.grad, rtol=5e-2, atol=1e-3 - ), "occ gradients don't match" - - def test_gradients_finite_differences(self, grad_tensors): - """Verify gradients match finite differences for a simple case.""" - data = grad_tensors(n_atoms=5, n_voxels=20, grid_shape=(8, 8, 8)) - - def func(A): - density_map = torch.zeros(data["grid_shape"], dtype=A.dtype) - result = vectorized_add_to_map( - data["surrounding_coords"], data["voxel_indices"], density_map, - data["xyz"].detach(), data["b"].detach(), - data["inv_frac_matrix"], data["frac_matrix"], - A, data["B"].detach(), data["occ"].detach() - ) - return result.sum() - - # Check gradients for A using finite differences - A_test = data["A"].detach().clone().requires_grad_(True) - eps = 1e-5 - - # Compute analytical gradient - loss = func(A_test) - loss.backward() - analytical_grad = A_test.grad.clone() - - # Compute numerical gradient for a few elements - for i in range(min(3, A_test.shape[0])): - for j in range(min(2, A_test.shape[1])): - A_plus = A_test.detach().clone() - A_minus = A_test.detach().clone() - A_plus[i, j] += eps - A_minus[i, j] -= eps - - numerical_grad = (func(A_plus) - func(A_minus)) / (2 * eps) - relative_error = abs( - numerical_grad - analytical_grad[i, j] - ) / (abs(numerical_grad) + 1e-8) - - assert relative_error < 0.01, ( - f"Gradient mismatch at [{i},{j}]: " - f"numerical={numerical_grad:.6f}, " - f"analytical={analytical_grad[i,j]:.6f}" - ) - - -class TestPerformance: - """Performance benchmarks.""" - - @pytest.fixture - def setup_tensors(self): - """Create test tensors for benchmarking.""" - def _setup(n_atoms=100, n_voxels=500, grid_shape=(64, 64, 64), - device="cpu", dtype=torch.float64, seed=42): - torch.manual_seed(seed) - - cell = torch.tensor( - [50.0, 60.0, 70.0, 90.0, 90.0, 90.0], dtype=dtype, device=device - ) - frac_matrix = torch.diag(cell[:3]) - inv_frac_matrix = torch.diag(1.0 / cell[:3]) - - xyz = torch.rand(n_atoms, 3, device=device, dtype=dtype) * 30 + 10 - b = torch.rand(n_atoms, device=device, dtype=dtype) * 40 + 10 - A = torch.rand(n_atoms, 5, device=device, dtype=dtype) * 5 + 1 - B = torch.rand(n_atoms, 5, device=device, dtype=dtype) * 10 + 2 - occ = torch.rand(n_atoms, device=device, dtype=dtype) * 0.5 + 0.5 - - surrounding_coords = ( - xyz[:, None, :] - + torch.randn(n_atoms, n_voxels, 3, device=device, dtype=dtype) * 2 - ) - voxel_indices = torch.randint( - 0, min(grid_shape), (n_atoms, n_voxels, 3), device=device - ) - - return { - "surrounding_coords": surrounding_coords, - "voxel_indices": voxel_indices, - "grid_shape": grid_shape, - "xyz": xyz, - "b": b, - "frac_matrix": frac_matrix, - "inv_frac_matrix": inv_frac_matrix, - "A": A, - "B": B, - "occ": occ, - } - return _setup - - @pytest.mark.slow - def test_performance_improvement(self, setup_tensors): - """Auto-compiled version should be faster after warmup.""" - data = setup_tensors( - n_atoms=1000, n_voxels=2000, grid_shape=(128, 128, 128) - ) - - n_warmup = 3 - n_trials = 10 - - # Warmup both - for _ in range(n_warmup): - map_orig = torch.zeros(data["grid_shape"], dtype=data["xyz"].dtype) - original_add_to_map( - data["surrounding_coords"], data["voxel_indices"], map_orig, - data["xyz"], data["b"], data["inv_frac_matrix"], - data["frac_matrix"], data["A"], data["B"], data["occ"] - ) - - map_compiled = torch.zeros(data["grid_shape"], dtype=data["xyz"].dtype) - vectorized_add_to_map( - data["surrounding_coords"], data["voxel_indices"], map_compiled, - data["xyz"], data["b"], data["inv_frac_matrix"], data["frac_matrix"], - data["A"], data["B"], data["occ"] - ) - - # Benchmark original - times_orig = [] - for _ in range(n_trials): - map_orig = torch.zeros(data["grid_shape"], dtype=data["xyz"].dtype) - t0 = time.perf_counter() - original_add_to_map( - data["surrounding_coords"], data["voxel_indices"], map_orig, - data["xyz"], data["b"], data["inv_frac_matrix"], - data["frac_matrix"], data["A"], data["B"], data["occ"] - ) - times_orig.append(time.perf_counter() - t0) - - # Benchmark auto-compiled - times_compiled = [] - for _ in range(n_trials): - map_compiled = torch.zeros(data["grid_shape"], dtype=data["xyz"].dtype) - t0 = time.perf_counter() - vectorized_add_to_map( - data["surrounding_coords"], data["voxel_indices"], map_compiled, - data["xyz"], data["b"], data["inv_frac_matrix"], data["frac_matrix"], - data["A"], data["B"], data["occ"] - ) - times_compiled.append(time.perf_counter() - t0) - - mean_orig = np.mean(times_orig) - mean_compiled = np.mean(times_compiled) - speedup = mean_orig / mean_compiled - - print(f"\nOriginal: {mean_orig*1000:.2f} ms") - print(f"Auto-compiled: {mean_compiled*1000:.2f} ms") - print(f"Speedup: {speedup:.2f}x") - - # Assert some speedup (at least not significantly slower) - assert speedup > 0.8, f"Compiled version too slow: {speedup:.2f}x" - - -class TestWarmup: - """Test warmup function.""" - - def test_warmup_no_error(self): - """warmup should not raise errors.""" - warmup(device="cpu") - - def test_warmup_enables_fast_execution(self): - """After warmup, function should work correctly.""" - # Warmup - warmup(device="cpu") - - # Create test data and verify function works - n_atoms, n_voxels = 10, 50 - grid_shape = (16, 16, 16) - dtype = torch.float32 - - surrounding_coords = torch.randn(n_atoms, n_voxels, 3, dtype=dtype) - voxel_indices = torch.randint(0, 16, (n_atoms, n_voxels, 3)) - density_map = torch.zeros(grid_shape, dtype=dtype) - xyz = torch.randn(n_atoms, 3, dtype=dtype) - b = torch.rand(n_atoms, dtype=dtype) * 50 + 10 - inv_frac_matrix = torch.eye(3, dtype=dtype) * 0.02 - frac_matrix = torch.eye(3, dtype=dtype) * 50 - A = torch.rand(n_atoms, 5, dtype=dtype) - B = torch.rand(n_atoms, 5, dtype=dtype) * 10 + 1 - occ = torch.ones(n_atoms, dtype=dtype) - - result = vectorized_add_to_map( - surrounding_coords, voxel_indices, density_map, xyz, b, - inv_frac_matrix, frac_matrix, A, B, occ - ) - - assert result.shape == grid_shape - assert torch.isfinite(result).all() diff --git a/tests/unit/io/test_data.py b/tests/unit/io/test_data.py index 2cbd4c51..a615bfba 100644 --- a/tests/unit/io/test_data.py +++ b/tests/unit/io/test_data.py @@ -80,14 +80,14 @@ def test_reflection_data_cpu(self): @pytest.mark.unit @pytest.mark.gpu - def test_reflection_data_cuda(self, gpu_device): - """Test CUDA movement.""" + def test_reflection_data_accelerator(self, gpu_device): + """Movement onto whichever accelerator this host has.""" from torchref.io import ReflectionData data = ReflectionData() - data = data.cuda() + data = data.to(gpu_device) - assert data.device.type == "cuda" + assert data.device.type == gpu_device.type class TestReflectionDataAttributes: diff --git a/tests/unit/io/test_reflection_data_reindex.py b/tests/unit/io/test_reflection_data_reindex.py new file mode 100644 index 00000000..691a5526 --- /dev/null +++ b/tests/unit/io/test_reflection_data_reindex.py @@ -0,0 +1,179 @@ +"""Regression tests for per-reflection field reindexing. + +``validate_hkl`` / ``remap`` / ``reduce_to_spacegroup`` must carry EVERY +per-reflection field onto the new HKL grid, not a hand-maintained subset. The +historical bug left ``hkl_anomalous`` (read by ``hkl_for_sf``) at the +pre-alignment length, which crashed difference refinement whenever the dark and +light datasets had different reflection sets. See docs/changelog.rst 0.6.2. +""" + +import pytest +import torch + +from torchref.io.datasets.reflection_data import ReflectionData + + +def _base_grid(h=10, k=10, lmax=10): + """A block of Miller indices with strictly-positive l (already in the ASU).""" + hs = torch.arange(-h, h + 1) + ks = torch.arange(-k, k + 1) + ls = torch.arange(1, lmax + 1) + return ( + torch.stack(torch.meshgrid(hs, ks, ls, indexing="ij"), dim=-1) + .reshape(-1, 3) + .to(torch.int32) + ) + + +def _synthetic(hkl, seed=0, device="cpu"): + n = hkl.shape[0] + g = torch.Generator().manual_seed(seed) + F = torch.rand(n, generator=g) * 100.0 + 1.0 + return ReflectionData.from_tensors( + hkl, + F, + F * 0.1, + (60.0, 60.0, 60.0, 90.0, 90.0, 90.0), + "P 21 21 21", + device=device, + verbose=0, + ) + + +class TestValidateHklReindex: + """The reported crash lives in validate_hkl (collection HKL alignment).""" + + def test_carries_all_per_reflection_fields(self): + grid = _base_grid() + n = grid.shape[0] + # ``light`` lacks the last 10%; the reference grid lacks the first 10%, + # so the two sets genuinely differ (each has reflections the other lacks). + light = _synthetic(grid[: int(n * 0.9)], seed=1) + ref_hkl = grid[int(n * 0.1):].clone() + assert len(light.hkl_anomalous) == len(light.hkl) # sane before + + light.validate_hkl(ref_hkl) + + m = len(light.hkl) + assert m == len(ref_hkl) + # The field the old code left stale (the direct cause of the crash): + assert light.hkl_anomalous.shape[0] == m + assert light.hkl_for_sf().shape[0] == m + # Derived-from-HKL fields recompute for the new grid: + assert light.centric.shape[0] == m + assert light.friedel_flags.shape[0] == m + # Global invariant: no per-reflection tensor left at the old length. + light._assert_per_reflection_consistent() + + def test_identical_hkl_preserves_count(self): + grid = _base_grid(6, 6, 6) + d = _synthetic(grid, seed=2) + n0 = len(d.hkl) + d.validate_hkl(d.hkl.clone()) + assert len(d.hkl) == n0 + d._assert_per_reflection_consistent() + + +class TestP1RoundTripReindex: + """The same class of bug lived latently in remap/expand_to_p1 and + reduce_to_spacegroup (silent data loss rather than a crash).""" + + def test_expand_to_p1_carries_validation_flags(self): + grid = _base_grid(6, 6, 6) + d = _synthetic(grid, seed=3) + d.generate_validation_set(val_fraction_of_free=0.5, seed=0) + assert d.validation_flags is not None + + p1 = d.expand_to_p1() + # validation_flags used to be dropped to None by remap; now carried. + assert p1.validation_flags is not None + assert p1.validation_flags.shape[0] == len(p1.hkl) + p1._assert_per_reflection_consistent() + + def test_reduce_to_spacegroup_consistent(self): + grid = _base_grid(6, 6, 6) + d = _synthetic(grid, seed=4) + p1 = d.expand_to_p1() + back = p1.reduce_to_spacegroup("P 21 21 21") + back._assert_per_reflection_consistent() + + +@pytest.mark.integration +class TestCollectionDifferenceMismatchedHKL: + """In-process reproduction of the reported failure: a DatasetCollection + built from two different reflection sets, run through the collection + difference target's forward (the exact call site that crashed).""" + + def test_forward_finite_with_different_reflection_sets(self, pdb_dir, mtz_dir): + pdb = pdb_dir / "3GR5.pdb" + mtz = mtz_dir / "3GR5.mtz" + if not (pdb.exists() and mtz.exists()): + pytest.skip("3GR5 fixture not present") + + from torchref.cli._common import load_model + from torchref.config import get_default_device + from torchref.io.datasets.collection import DatasetCollection + from torchref.io.datasets.reflection_data import ReflectionData + from torchref.model.model_collection import ModelCollection + from torchref.refinement.targets.collection import CollectionDifferenceTarget + from torchref.scaling.collection_scaler import CollectionScaler + + dev = get_default_device() + full = ReflectionData(device=dev, verbose=0).load_mtz(str(mtz)) + n = len(full.hkl) + idx = torch.arange(n, device=full.hkl.device) + # Two DIFFERENT reflection sets with DIFFERENT counts (drops the last + # 10% vs the first 15%) -> reproduces the shape-mismatch crash pre-fix. + dark = full.__select__(idx < int(n * 0.90)) + light = full.__select__(idx >= int(n * 0.15)) + assert len(dark.hkl) != len(light.hkl) + + dc = DatasetCollection(device=dev, verbose=0) + dc.add_dataset("dark", dark, set_as_reference=True) + dc.add_dataset("light", light) # -> validate_hkl aligns onto dark grid + + # High-resolution limit for the model FFT grid (covers the data). + d_min = float(full.resolution.min()) + model_dark = load_model(str(pdb), max_res=d_min, device=dev, verbose=0) + model_light = load_model(str(pdb), max_res=d_min, device=dev, verbose=0) + mc = ModelCollection([model_dark, model_light], dark_key="dark", verbose=0) + mc.add_dark() + mc.add_timepoint("light", [0.7, 0.3]) + + scaler = CollectionScaler(dc, mc, verbose=0) + scaler.initialize() + + target = CollectionDifferenceTarget(dc, mc, scaler=scaler, verbose=0) + loss = target.forward() + assert torch.isfinite(loss) + + +@pytest.mark.unit +def test_reindex_preserves_non_per_reflection_u_aniso(): + """``U_aniso`` must be exempt by *name*, not by a shape coincidence. + + The reindexer decides "is this per-reflection?" by ``shape[0] == n_hkl``. + That heuristic collides when the dataset happens to have exactly as many + reflections as the field is long -- ``U_aniso`` is ``(6,)``, so a + 6-reflection dataset would have it gathered and reordered as if it were + per-reflection data. + """ + # Exactly 6 reflections: the same length as U_aniso. + hkl = torch.tensor( + [[1, 0, 1], [0, 1, 1], [0, 0, 1], [1, 1, 1], [1, 0, 2], [0, 1, 2]], + dtype=torch.int32, + ) + data = _synthetic(hkl) + u_aniso = torch.tensor([0.1, 0.2, 0.3, 0.01, 0.02, 0.03]) + data.U_aniso = u_aniso.clone().to(device=data.device) + assert len(data.hkl) == data.U_aniso.shape[0] == 6, "precondition: lengths collide" + + # Reindex onto a different HKL ordering/size; U_aniso must not follow. + ref_hkl = torch.tensor( + [[0, 1, 2], [1, 0, 2], [1, 1, 1], [0, 0, 1], [0, 1, 1], [1, 0, 1], [2, 0, 1]], + dtype=torch.int32, + ) + data.validate_hkl(ref_hkl.to(data.device)) + + assert data.U_aniso.shape == (6,), f"U_aniso reshaped to {tuple(data.U_aniso.shape)}" + assert torch.allclose(data.U_aniso.cpu(), u_aniso), "U_aniso values were permuted" diff --git a/tests/unit/model/test_model.py b/tests/unit/model/test_model.py index 7f584546..1781a89a 100644 --- a/tests/unit/model/test_model.py +++ b/tests/unit/model/test_model.py @@ -101,7 +101,7 @@ def test_model_gpu_device(self, gpu_device): model = Model(device=gpu_device) - assert model.device.type == 'cuda' + assert model.device.type == gpu_device.type class TestModelGetSelectionMask: diff --git a/tests/unit/model/test_parameter_wrappers.py b/tests/unit/model/test_parameter_wrappers.py index bb1b6b07..5cfeb037 100644 --- a/tests/unit/model/test_parameter_wrappers.py +++ b/tests/unit/model/test_parameter_wrappers.py @@ -133,7 +133,7 @@ def test_mixed_tensor_gpu(self, random_coordinates, gpu_device): mixed = MixedTensor(values, device=gpu_device) result = mixed() - assert result.device.type == 'cuda' + assert result.device.type == gpu_device.type class TestMixedTensorSet: diff --git a/tests/unit/refinement/test_estimate_beta_determinism.py b/tests/unit/refinement/test_estimate_beta_determinism.py index 4b84bdcb..ef1f7a4b 100644 --- a/tests/unit/refinement/test_estimate_beta_determinism.py +++ b/tests/unit/refinement/test_estimate_beta_determinism.py @@ -47,7 +47,7 @@ def _hard_inputs(dtype=torch.float32): @pytest.mark.unit -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.cuda class TestEstimateBetaGPUDeterminism: def test_gpu_bit_identical_across_calls(self): """Same input -> identical beta on every GPU call (was: 15x spread).""" diff --git a/tests/unit/refinement/test_forcefield_target.py b/tests/unit/refinement/test_forcefield_target.py index dacebdf3..df5db224 100644 --- a/tests/unit/refinement/test_forcefield_target.py +++ b/tests/unit/refinement/test_forcefield_target.py @@ -35,6 +35,21 @@ def __init__(self, n_atoms=5, has_hydrogens=True): # Create coordinates as a parameter for gradient testing self._coords = nn.Parameter(torch.randn(len(self.Z), 3)) + @property + def device(self): + """Device of the model's coordinates. + + Real ``Model`` / ``ModelFT`` expose this, and ``ModelTarget`` now + reconciles its own device against the model's, so a stand-in has to + carry it too. + """ + return self._coords.device + + @property + def dtype_float(self): + """Float dtype of the model's coordinates (mirrors ``Model``).""" + return self._coords.dtype + def xyz(self): """Return current coordinates.""" return self._coords diff --git a/tests/unit/refinement/test_ml_sigmaa.py b/tests/unit/refinement/test_ml_sigmaa.py index a264d2da..3aa8e556 100644 --- a/tests/unit/refinement/test_ml_sigmaa.py +++ b/tests/unit/refinement/test_ml_sigmaa.py @@ -189,15 +189,20 @@ def test_beta_physical_and_falls_with_resolution(self, refinement): assert (beta[v] > 0).all() and torch.isfinite(beta[v]).all() # The falling-with-resolution trend is an estimator-math property, - # checked in float64 so it is device-independent. + # checked in float64 so it is device-independent. Move to CPU before + # casting -- MPS has no float64, so the double() must happen off-device. data = refinement.reflection_data with torch.no_grad(): - fc = torch.abs(t._scaled_F_calc_full()).double().reshape(-1) - fo = data.get_corrected_data()[0].double().reshape(-1) - epsd = epsilon_from_hkl(data.hkl, getattr(data, "spacegroup", None)).double() + fc = torch.abs(t._scaled_F_calc_full()).cpu().double().reshape(-1) + fo = data.get_corrected_data()[0].cpu().double().reshape(-1) + epsd = epsilon_from_hkl( + data.hkl, getattr(data, "spacegroup", None) + ).cpu().double() s = get_scattering_vectors(data.hkl, data.cell) - dss = (torch.norm(s, dim=1) ** 2).double() - _b, bbin, _ = estimate_beta(fo, fc, data.centric, epsd, dss, data.free.mask) + dss = (torch.norm(s, dim=1) ** 2).cpu().double() + _b, bbin, _ = estimate_beta( + fo, fc, data.centric.cpu(), epsd, dss, data.free.mask.cpu() + ) assert (bbin > 0).all() and torch.isfinite(bbin).all() # beta is an absolute model-error variance in F^2 units (~(1-sigma_A^2)* # Sigma_N); Sigma_N ~= decays steeply with resolution, so absolute @@ -216,3 +221,42 @@ def test_gradient_reaches_model(self, refinement): refinement.xray_target_work.forward().backward() xyz = refinement.xray_target_work._model.xyz.refinable_params assert xyz.grad is not None and torch.isfinite(xyz.grad).all() + + +@pytest.mark.unit +def test_epsilon_from_hkl_returns_on_hkl_device(): + """``epsilon_from_hkl`` must answer on ``hkl``'s device, not the spacegroup's. + + ``SpaceGroup.apply_to_hkl`` moves its input onto the symmetry matrices' + device, so a spacegroup built elsewhere than the reflection data would + otherwise either raise inside the ``Hs == h0`` comparison or hand back a + tensor the caller cannot multiply against its per-reflection data. + """ + from torchref.base.targets.xray_ml_sigmaa import epsilon_from_hkl + from torchref.symmetry import SpaceGroup + + hkl = torch.tensor([[1, 0, 0], [0, 2, 0], [1, 1, 1]], dtype=torch.int32) + sg = SpaceGroup("P 21 21 21", device="cpu") + + eps = epsilon_from_hkl(hkl, sg) + assert eps.device == hkl.device + assert eps.shape == (3,) + assert (eps >= 1).all() + + # No spacegroup: the early-return path must honour the same contract. + assert epsilon_from_hkl(hkl, None).device == hkl.device + + +@pytest.mark.mps +def test_epsilon_from_hkl_cross_device(): + """hkl on the accelerator, spacegroup on CPU: still answers on hkl's device.""" + from torchref.base.targets.xray_ml_sigmaa import epsilon_from_hkl + from torchref.symmetry import SpaceGroup + + hkl = torch.tensor( + [[1, 0, 0], [0, 2, 0], [1, 1, 1]], dtype=torch.int32, device="mps" + ) + sg = SpaceGroup("P 21 21 21", device="cpu") + eps = epsilon_from_hkl(hkl, sg) + assert eps.device.type == "mps" + assert (eps.cpu() >= 1).all() diff --git a/tests/unit/refinement/test_targets.py b/tests/unit/refinement/test_targets.py index e8d87944..fe2728e7 100644 --- a/tests/unit/refinement/test_targets.py +++ b/tests/unit/refinement/test_targets.py @@ -160,7 +160,7 @@ def test_target_gpu_tensors(self, mock_F_obs, mock_F_sigma, gpu_device): loss = torch.mean((fobs / sigma) ** 2) - assert loss.device.type == 'cuda' + assert loss.device.type == gpu_device.type assert torch.isfinite(loss) diff --git a/tests/unit/restraints/test_restraints.py b/tests/unit/restraints/test_restraints.py index a4ba9d3a..1ec929cf 100644 --- a/tests/unit/restraints/test_restraints.py +++ b/tests/unit/restraints/test_restraints.py @@ -176,7 +176,7 @@ def test_bond_distance_gpu(self, random_coordinates, gpu_device): distance = torch.sqrt(torch.sum((coords[0] - coords[1]) ** 2)) - assert distance.device.type == 'cuda' + assert distance.device.type == gpu_device.type class TestRestraintNumericStability: diff --git a/tests/unit/scaling/test_scaler.py b/tests/unit/scaling/test_scaler.py index 77d5aa6a..44353856 100644 --- a/tests/unit/scaling/test_scaler.py +++ b/tests/unit/scaling/test_scaler.py @@ -81,7 +81,7 @@ def test_scaler_gpu_device(self, gpu_device): scaler = Scaler(device=gpu_device) - assert scaler.device.type == 'cuda' + assert scaler.device.type == gpu_device.type class TestScalingCalculations: diff --git a/tests/unit/scattering/test_anomalous.py b/tests/unit/scattering/test_anomalous.py index 191e0dc3..145f9b3a 100644 --- a/tests/unit/scattering/test_anomalous.py +++ b/tests/unit/scattering/test_anomalous.py @@ -293,18 +293,65 @@ def test_friedel_pair_asymmetry(self, test_pdb_file): sf_plus = model.get_structure_factor(hkl, apply_anomalous=True, recalc=True) sf_minus = model.get_structure_factor(-hkl, apply_anomalous=True, recalc=True) - # With anomalous scattering, F(h) != F(-h)* (not complex conjugates) - # due to the imaginary f'' contribution - # Note: The difference might be small, but should exist - is_conjugate = torch.allclose(sf_plus, sf_minus.conj(), atol=1e-6) - - # If Fe contributes significant anomalous scattering, they should not be conjugates - # We check this indirectly - the cache should have anomalous scatterers + # Friedel's law is governed by f'', which is imaginary and does not change sign + # with h, so it makes F(h) != F(-h)*. f' is real and dispersive and leaves the + # conjugate relation intact. + # + # In this model f'' is gated on ``apply_bijvoet`` (``ModelFT.__init__``, applied at + # ``model_ft.py:951`` via ``include_fdp``), which defaults to False because merged + # data is the usual target and Friedel-preserving F is correct for it. So the + # default path deliberately does *not* break Friedel's law -- both branches are + # asserted here rather than only the one this test originally assumed. + # + # History: this test previously computed ``is_conjugate`` and then ended in + # ``pass``, asserting nothing. A first attempt to fix it asserted breakdown on the + # default path and failed, because that path is Friedel-preserving by design. mask, _, _, _, _ = model._get_anomalous_cache() - if mask.any(): - # There are anomalous scatterers, so Friedel pairs should differ - # (in principle - the actual difference depends on the positions) - pass # The test passes if no exception was raised + assert mask.any(), ( + "no anomalous scatterers in this structure, so neither branch below is " + "meaningful -- pick a structure with an anomalous element" + ) + + def asymmetry(fp, fm): + return float( + (fp - fm.conj()).abs().norm().detach() / fp.abs().norm().detach() + ) + + # --- default branch: f' only, Friedel's law preserved ------------------- + sf_plus_no = model.get_structure_factor(hkl, apply_anomalous=False, recalc=True) + floor = asymmetry(sf_plus_no, model.get_structure_factor( + -hkl, apply_anomalous=False, recalc=True)) + default_asym = asymmetry(sf_plus, sf_minus) + + assert not bool(model.anomalous_bijvoet), "fixture unexpectedly enabled f''" + assert default_asym < 10.0 * max(floor, 1e-9), ( + f"the default path broke Friedel's law ({default_asym:.3e} against a " + f"no-anomalous floor of {floor:.3e}). With apply_bijvoet=False, f'' is zeroed " + "and F(h) must stay conjugate to F(-h)" + ) + # ...but f' must still be reaching F, or "anomalous" is doing nothing at all. + dispersive = float( + (sf_plus - sf_plus_no).abs().norm().detach() / sf_plus_no.abs().norm().detach() + ) + assert dispersive > 1e-4, ( + f"apply_anomalous=True changed F by only {dispersive:.3e}, so the dispersive " + "f' term is not reaching the structure factors either" + ) + + # --- apply_bijvoet=True: f'' included, Friedel's law broken ------------- + bij = ModelFT( + wavelength=1.0, anomalous_threshold=0.5, apply_bijvoet=True, verbose=0 + ) + bij.load_pdb(test_pdb_file) + bij_asym = asymmetry( + bij.get_structure_factor(hkl, apply_anomalous=True, recalc=True), + bij.get_structure_factor(-hkl, apply_anomalous=True, recalc=True), + ) + assert bij_asym > 100.0 * max(floor, 1e-9), ( + f"with apply_bijvoet=True the Friedel asymmetry is only {bij_asym:.3e} " + f"against a floor of {floor:.3e}, so f'' is not reaching the structure " + "factors and Bijvoet differences would be absent" + ) def test_gradient_flow(self, test_pdb_file): """Test that gradients flow through anomalous correction.""" diff --git a/tests/unit/structure_factor/__init__.py b/tests/unit/structure_factor/__init__.py new file mode 100644 index 00000000..851a3f20 --- /dev/null +++ b/tests/unit/structure_factor/__init__.py @@ -0,0 +1,301 @@ +"""Structure-factor correctness, gated against an oracle at every derivative order. + +Every accuracy gate in this package traces back to a reference that is *independent* +of the code under test. That is the whole point: before this package existed, the +density and SF coverage was + +* kernel-A-vs-kernel-B parity at ``rel < 2e-2`` (``test_variable_radius_{gpu,mps}.py``), + which cannot detect an error the two kernels share; +* splat-vs-brute-force-splat (``test_canonical_sphere_cpu.py``), which validates the + truncation *contract* rather than the physics; +* finite differences of the map route against itself, which is self-consistency. + +Nothing checked that the FFT route reproduces an analytic ``F(hkl)``, and nothing +checked its absolute scale at all -- every gate was ratio-, cosine- or parity-based, so +a factor-of-volume slip passed the entire suite. + +The chain +--------- +:: + + gemmi -> direct summation -> the FFT/splat route (amplitudes) + finite differences -> direct summation -> the FFT/splat route (grad, HVP) + +**Direct summation is the oracle.** It evaluates ``F(h) = sum_j occ_j f_j(s) DW_j +exp(2 pi i h.x_j)`` analytically -- no grid, no truncation radius -- so finite differences +are a valid reference for it (measured 1.19e-10). It is then an *independent* reference for +the map route, which is the part finite differences cannot supply: FD differentiates the +same discretized function the kernel does, so it confirms only that autograd is correct, +never that the kernel is. Measured, the map route's HVP agrees with itself to 5e-04 while +sitting 2.3e-02 from the analytic answer -- see +``test_second_order.py::test_finite_differences_cannot_detect_map_route_error``. + +**gemmi anchors the forward values.** It is a separate codebase, is a pip dependency +(so it runs in CI, unlike conda-only cctbx), and exposes no derivatives -- hence FD +remains the derivative link. + +What the gemmi comparison does and does not prove +------------------------------------------------- +torchref's ITC92 table was *generated from gemmi* +(``torchref/scripts/generate_scattering_table.py`` writes +``torchref/data/itc92_scattering_factors.pt``). Verified here: torchref stores gemmi's +four ITC92 Gaussians plus the constant ``c`` folded in as a fifth with ``B = 0``, and +``f(0)`` agrees to every printed digit. + +So ``f(s)`` is a **shared** input and cancels out of a DS-vs-gemmi comparison. That test +does **not** validate the form factors. What it does validate is everything layered on +top of them, which is where the error-prone conventions live: the ``2 pi`` factors, the +phase sign, Debye-Waller in both the B and U conventions, the +``[U11,U22,U33,U12,U13,U23]`` ordering, occupancy weighting, and -- in the non-P1 case -- +the symmetry algebra against a second implementation. + +Two consequences, both acted on: + +1. **The gate is tight, not loose.** With the coefficients shared there is no table + discrepancy to absorb. See the measured numbers on ``RTOL_VS_GEMMI`` below. +2. **The table needs its own check**, since the SF comparison cannot provide one. + ``test_forward.py::test_stored_table_matches_gemmi`` compares the stored ``.pt`` + against gemmi's live table element by element -- the one place gemmi *is* an + independent reference for ``f(s)``, and the only thing that would catch table drift, + a corrupted regeneration, or a gemmi release that changed coefficients. + +Why the oracle is ``_eager_*`` and never ``ds_iso``/``SfDS`` +------------------------------------------------------------ +``ds_iso``, ``ds_aniso`` and ``SfDS`` all route through ``_CheckpointedSF``, whose +``backward`` calls ``torch.autograd.grad`` **without** ``create_graph=True`` on detached +copies (``torchref/base/direct_summation/dispatch.py:197``). They are exact at first +order -- measured agreement with the eager path is 2.3e-16 -- but a second derivative +through them raises ``element 0 of tensors does not require grad``. ``force_portable`` +does not help; it only steers away from Triton and still lands on ``_CheckpointedSF``. + +So the oracle is ``_eager_iso`` / ``_eager_aniso``, which are pure torch and therefore +compose to any order. ``test_second_order.py`` asserts both facts -- that the eager path +``gradgradcheck``s and that the public API raises -- so the constraint is documented +behaviour rather than a trap. +""" + +from __future__ import annotations + +# --------------------------------------------------------------------------- +# Outer link: gemmi -> DS. Forward only; gemmi exposes no derivatives. +# --------------------------------------------------------------------------- +# Measured on first calibration (float64, CPU), rel L2 on complex F: +# +# iso, P1, 3GR5 (1329 atoms, 240 refl) 3.28e-08 +# aniso, P1, 7L84 (1209 ANISOU atoms, 200 refl) 2.91e-08 +# iso, P 65 2 2, 3GR5 via SfDS (12 ops, 200 refl) 3.22e-08 +# +# The floor is the **stored table's float32 precision**. ``torchref/data/ +# itc92_scattering_factors.pt`` holds float32 ``(104, 5)`` tensors, while gemmi's Python +# ``it92`` returns doubles; the worst relative coefficient difference over Z=1..103 is +# 5.9e-08, which is half of float32 eps. Asking ``get_scattering_params_by_z`` for +# float64 widens those float32 values rather than recovering gemmi's doubles, so even the +# "float64" oracle carries ~6e-8 relative error in ``f(s)``. +# +# That is the correct trade for a float32 production path, not a defect -- but it means +# tightening below ~1e-7 would gate on table storage precision rather than on any +# structure-factor logic. ``test_stored_table_matches_gemmi`` pins the rounding exactly +# (bit-exact in float32) so this floor cannot drift unnoticed. +# +# Non-vacuity, measured with the same harness: permuting the ADP off-diagonals to +# swap U12<->U13 moves rel L2 to 1.35e-02, and reversing them to 6.08e-03 -- five +# orders of magnitude clear of the gate. The comparison genuinely constrains the +# convention. +RTOL_VS_GEMMI = 1e-6 # rel L2 on complex F; 30x margin on the worst measurement +MAXREL_VS_GEMMI = 1e-4 # max per-reflection; 48x margin (worst was 2.06e-06) +SCALE_TOL_GEMMI = 1e-6 # |least-squares scale - 1|; worst measured 1.3e-08 + +# Stored table vs gemmi's live table. Both hold the same float32-rounded coefficients, +# so this is expected to be exact; a nonzero result is a finding, not a tolerance to +# widen. +ATOL_TABLE_VS_GEMMI = 0.0 + +# --------------------------------------------------------------------------- +# Oracle soundness: FD -> DS. DS is analytic and smooth, so gate hard. +# --------------------------------------------------------------------------- +EPS_GRADCHECK = 1e-6 +ATOL_GRADCHECK = 1e-5 +RTOL_GRADCHECK = 1e-3 +RTOL_DS_HVP_VS_FD = 1e-5 # measured 1.19e-10 on a 6-atom scene + +# --------------------------------------------------------------------------- +# The FFT/splat route vs the DS oracle. +# --------------------------------------------------------------------------- +# These gates bound a real discretization error rather than eliminating it: the map route +# carries both a truncation and a grid-sampling residual, and gradients are markedly more +# sensitive to both than amplitudes are, because each xyz derivative brings a factor +# ~``2*pi*h`` and so re-weights the comparison toward high-resolution reflections -- the +# ones the grid resolves worst. +# +# **Production runs in float32.** So float32 is the load-bearing case throughout this +# package, not a variant bolted on beside float64 -- every parametrization lists it +# first, and these gates are calibrated on it. float64 is kept because it separates a +# genuine truncation residual from float32 accumulation noise: if a float32 result misses +# a gate and float64 passes it comfortably, the cause is precision, not the kernel. The +# oracle itself always stays in float64 regardless of the candidate's dtype -- a +# reference should not inherit the precision of the thing it is judging. +# Measured at the **production** grid (``max_res = d_min``, spacing ``d_min/3``) with the +# default 3.0 sigma cutoff, differentiating ``ls_target`` against pseudo-observations +# offset from the oracle by 10% relative noise. Identical to 3 significant figures across +# float32/float64 and across both CPU kernels, so these are properties of the +# discretization, not of precision or of a kernel: +# +# amplitude g_xyz g_occ g_adp/g_U HVP HVP cos +# iso 5.89e-03 8.51e-02 2.68e-02 4.45e-02 2.26e-02 0.99980 +# aniso 5.20e-03 8.17e-02 1.20e-01 2.24e-01 1.39e-02 0.99991 +# +# The anisotropic ADP gradient is the outlier at 2.2e-01 -- it converges to 5.6e-03 by +# fineness 1.6, so it is grid-sampling error rather than a defect in the U derivative, but +# it does not fit inside 1e-1 at production sampling. Rather than raise the shared gate +# (which would stop it constraining anything) or silently test on a finer grid than +# production, aniso ADP gets its own documented constant. +# **Authoritative gates are calibrated on a real structure, not a synthetic scene.** +# Measured on 7L84 forced to P1 -- 1209 atoms all carrying ANISOU, 799 reflections to +# 1.5 A -- at production oversampling 3, float32 candidate against the float64 oracle: +# +# amplitude rel L2 2.28e-03 +# xyz gradient 1.35e-02 cos 0.999909 +# U gradient 2.62e-02 cos 0.999657 gradnorm ratio 0.9997 +# +# Synthetic small-cell scenes are markedly *pessimistic* here and must not be used to set +# absolute gates. Aliasing contamination in the derivative cancels across atoms roughly as +# 1/sqrt(N), because with few atoms each reflection's F is dominated by a handful of +# coherent contributions while with many the inter-atom phases are quasi-random. Measured +# deviatoric-U gradnorm ratio at production oversampling, all else equal: +# +# 10 atoms 1.43 - 1.70 (independent of how anisotropic U actually is) +# 60 atoms 0.972 +# 300 atoms 0.997 +# 7L84, 1209 0.9995 +# +# An earlier revision of this package calibrated the aniso gates on a 10-atom scene and +# recorded a "systematic 54% deviatoric ADP gradient bias" as a production finding. It was +# an atom-count artifact and does not exist on real structures. The synthetic scenes are +# retained for cheap parametrized coverage -- dtype, engine, cell angle -- where the +# comparison is *between* backends and scene realism is irrelevant. +RTOL_AMPLITUDE = 1e-2 # 7L84 2.28e-03; synthetic worst 5.50e-03 +RTOL_GRADIENT = 1e-1 # 7L84 worst leaf 2.62e-02 +RTOL_HVP = 1e-1 # synthetic worst 2.22e-02 +COS_MIN = 0.999 # forward complex F +COS_MIN_GRADIENT = 0.999 # 7L84 worst 0.999657 +SCALE_TOL = 1e-3 # |best-fit scale - 1|; worst measured 3.7e-05 + +# Synthetic-scene coverage sweeps only, never as an accuracy statement about production. +# Small cells with few atoms overstate the discretization residual and which leaf is worst +# moves around with atom count (10 atoms: aniso U; 60 atoms: iso adp at 1.36e-01). +RTOL_GRADIENT_SYNTHETIC = 2e-1 +COS_MIN_GRADIENT_SYNTHETIC = 0.99 + +# --------------------------------------------------------------------------- +# Cross-backend parity. Same truncation contract on every path, so near-exact. +# --------------------------------------------------------------------------- +RTOL_BACKEND_F32 = 2e-4 +RTOL_BACKEND_F64 = 1e-12 + +# Gradients need their own float32 constant. The fused C++ kernel's hand-written backward +# and the portable splat's autograd accumulate in different orders, and float32 does not +# forgive that the way the forward pass does. Measured fused-vs-portable on ``scene_fine`` +# (60 atoms), float32: +# +# xyz occ adp/U +# iso 2.07e-04 7.63e-04 1.43e-03 +# aniso 1.95e-04 6.79e-04 6.29e-04 +# +# so ~7x above the 2e-04 that suffices for forward values, worst case on the iso ADP +# gradient. float64 stays exact to 1e-12, which is the evidence that this is accumulation +# order rather than a genuine difference in what the two kernels compute. +# +# Unlike the accuracy gates, this one is *comparative* -- two backends on identical inputs +# -- so it is legitimately calibrated on a synthetic scene: the atom-count effect that +# makes small scenes unrepresentative for DS comparisons cancels when both sides share it. +RTOL_BACKEND_GRAD_F32 = 5e-3 +RTOL_BACKEND_GRAD_F64 = 1e-10 + +# --------------------------------------------------------------------------- +# Direct-summation kernels vs the eager oracle. +# --------------------------------------------------------------------------- +# Distinct from the map-route gates above: these compare one analytic implementation +# against another, with no grid and no truncation in between, so the only source of +# disagreement is float arithmetic. They are gated near precision accordingly. +# +# Measured on the 60-atom scene, candidate vs the CPU float64 ``_eager_*`` oracle: +# +# device dtype kernel F rel L2 worst grad rel L2 grad cos +# cpu float32 checkpointed 1.71e-06 2.39e-05 1.0000000 +# cpu float64 checkpointed 2.44e-16 2.69e-15 1.0000000 +# mps float32 checkpointed 1.72e-06 2.41e-05 1.0000000 +# +# The MPS row is a path that had no coverage of any kind before: there is no Metal +# direct-summation kernel, so any MPS device lands on +# ``_checkpointed_*`` running on-device. It agrees with the CPU float32 row to the digit. +# +# ``ds_triton`` is **unmeasured** -- no CUDA on the calibration host. It shares these gates, +# which is a prediction. The deleted ``test_ds_triton_vs_eager.py`` used a max-abs-diff +# gate of 1e-3 on a different metric; if a Triton leg misses 1e-4 here, measure before +# widening. +RTOL_DS_F32 = 1e-4 # worst measured 2.4e-05 +RTOL_DS_F64 = 1e-13 # worst measured 2.7e-15 +COS_MIN_DS = 0.9999 + +# --------------------------------------------------------------------------- +# Accelerators need no constants of their own. Measured, not assumed. +# --------------------------------------------------------------------------- +# The device legs were expected to need looser gates than the CPU ones, because an +# accelerator adds its own kernel arithmetic on top of the shared discretization error. +# They do not. Measured on the 60-atom synthetic scene against the CPU float64 DS oracle, +# iso then aniso, worst case per column: +# +# device dtype kernel amplitude g_xyz g_occ g_U/adp HVP +# cpu float32 cpu_sphere 4.37e-03 6.97e-02 3.61e-02 8.02e-02 1.62e-02 +# cpu float64 cpu_sphere 4.37e-03 6.98e-02 3.63e-02 8.09e-02 1.62e-02 +# cpu float64 portable 4.37e-03 6.98e-02 3.63e-02 8.09e-02 1.62e-02 +# mps float32 portable 4.37e-03 6.98e-02 3.63e-02 8.09e-02 1.62e-02 +# mps float32 mps_metal 4.37e-03 6.98e-02 3.63e-02 8.09e-02 raises +# +# Every row agrees to the printed digits. That is the truncation-contract standardization +# showing up as a measurement: with all four kernels applying the same atom-centred +# Cartesian sphere at the same raw policy radius, what is left over is *shared* +# discretization error, not per-kernel divergence. So the CPU constants above bound the +# accelerators, and adding ``RTOL_*_MPS`` / ``RTOL_*_CUDA`` would have created numbers with +# nothing behind them. +# +# ``mps_metal`` raising on the HVP is the documented contract, not a gap -- see +# ``mps/variable_radius.py:12``. ``test_second_order.py`` asserts both halves: which +# kernels give a correct second derivative, and which raise rather than returning a wrong +# one. +# +# **CUDA is unmeasured.** This host has no CUDA, so the Triton legs are written and marked +# but have never run. They share the parametrization and the tolerances with MPS, which is +# defensible given every other backend landed on the same numbers -- but it is a prediction +# until someone runs: +# +# pytest tests/unit/structure_factor -v -s -k cuda +# +# on a CUDA host. If a Triton leg misses a gate there, establish whether it is kernel +# arithmetic or a real geometry difference before touching a constant. + +__all__ = [ + "RTOL_VS_GEMMI", + "MAXREL_VS_GEMMI", + "SCALE_TOL_GEMMI", + "ATOL_TABLE_VS_GEMMI", + "EPS_GRADCHECK", + "ATOL_GRADCHECK", + "RTOL_GRADCHECK", + "RTOL_DS_HVP_VS_FD", + "RTOL_AMPLITUDE", + "RTOL_GRADIENT", + "RTOL_GRADIENT_SYNTHETIC", + "RTOL_HVP", + "COS_MIN", + "COS_MIN_GRADIENT", + "COS_MIN_GRADIENT_SYNTHETIC", + "SCALE_TOL", + "RTOL_BACKEND_F32", + "RTOL_BACKEND_F64", + "RTOL_BACKEND_GRAD_F32", + "RTOL_BACKEND_GRAD_F64", + "RTOL_DS_F32", + "RTOL_DS_F64", + "COS_MIN_DS", +] diff --git a/tests/unit/structure_factor/conftest.py b/tests/unit/structure_factor/conftest.py new file mode 100644 index 00000000..bcbf9be4 --- /dev/null +++ b/tests/unit/structure_factor/conftest.py @@ -0,0 +1,293 @@ +"""Fixtures for the structure-factor oracle package. + +Package-scoped fixtures capture the real redundancy -- three test modules otherwise +recomputing the same oracle -- while being recomputed on every run, so there is no +invalidation logic and no staleness risk at all. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch + +import torchref +from torchref.config import device as device_cfg, dtypes + +from tests.conftest import _accelerator + +from . import helpers as H + + +# --------------------------------------------------------------------------- +# Device axis +# --------------------------------------------------------------------------- +# Built at **import time**, copying ``tests/conftest.py:295``. That is load-bearing: the +# backend mark has to be attached during *collection*, because +# ``pytest_collection_modifyitems`` is what gates on it and cannot see a mark added later +# from inside a fixture. On a CPU-only host the accelerator param does not exist at all, +# so there is no skip noise -- and on this host the ``cuda`` leg likewise never appears. +# +# The accelerator carries its *specific* backend mark (``cuda`` or ``mps``), not the +# generic ``gpu`` one, so a CUDA-less host skips the cuda leg with an accurate reason. +_DEVICES = [pytest.param(torch.device("cpu"), id="cpu")] +_ACCELERATOR = _accelerator() +if _ACCELERATOR is not None: + _DEVICES.append( + pytest.param( + _ACCELERATOR, + id=_ACCELERATOR.type, + marks=getattr(pytest.mark, _ACCELERATOR.type), + ) + ) + +#: float32 first: the production dtype. MPS cannot hold float64 at all, which +#: ``helpers.device_supports_dtype`` filters -- so ``(mps, float64)`` yields no kernels +#: and therefore no test, rather than a test that skips or silently passes. +_DTYPES = [torch.float32, torch.float64] + + +def device_dtype_kernels(): + """Every ``(device, dtype, kernel_name)`` that names a real production path. + + Enumerated from the kernel registry rather than written out, so a kernel added to + ``helpers._KERNEL_SPECS`` is covered automatically and an unsupported combination + produces no test instead of a vacuous one. + """ + out = [] + for dev_param in _DEVICES: + device = dev_param.values[0] + for dtype in _DTYPES: + for name in H.kernels_for(device, dtype): + out.append( + pytest.param( + device, + dtype, + name, + id=f"{device.type}-{str(dtype).replace('torch.float', 'f')}-{name}", + marks=dev_param.marks, + ) + ) + return out + + +def ds_device_dtype_kernels(): + """Every ``(device, dtype, ds_kernel_name)`` naming a real direct-summation path. + + Separate from the splat list because the two families have different device/dtype + envelopes -- notably there is no Metal DS kernel, so the MPS leg here is + ``_checkpointed_*`` running on-device rather than a native shader. + """ + out = [] + for dev_param in _DEVICES: + device = dev_param.values[0] + for dtype in _DTYPES: + for name in H.ds_kernels_for(device, dtype): + out.append( + pytest.param( + device, dtype, name, + id=f"{device.type}-{str(dtype).replace('torch.float', 'f')}-{name}", + marks=dev_param.marks, + ) + ) + return out + + +DEVICE_DTYPE_KERNELS = device_dtype_kernels() +DS_DEVICE_DTYPE_KERNELS = ds_device_dtype_kernels() + + +# --------------------------------------------------------------------------- +# Global config +# --------------------------------------------------------------------------- +@pytest.fixture(scope="package", autouse=True) +def _float64_cpu(): + """float64/complex128 on CPU for this package; restore afterwards. + + Required, not cosmetic: ``iso_structure_factor_torched`` casts ``hkl`` to the + *global* ``dtypes.float`` (``torchref/base/direct_summation/isotropic.py:121``), so + under the default float32 config a float64 leaf produces a dtype-mismatched matmul. + That is why the pre-existing tests wrapped every eager-SF call in a ``double_cpu`` + fixture. + + ``sigma_cutoff_ed`` is restored here too -- the three copies of ``double_cpu`` this + replaces did not, so a test that changed the cutoff leaked it into everything that + ran after it. + """ + f0, c0, d0 = dtypes.float, dtypes.complex, device_cfg.current + s0 = torchref.sigma_cutoff_ed.value + dtypes.float = torch.float64 + dtypes.complex = torch.complex128 + device_cfg.current = torch.device("cpu") + try: + yield + finally: + dtypes.float = f0 + dtypes.complex = c0 + device_cfg.current = d0 + torchref.sigma_cutoff_ed.value = s0 + + +@pytest.fixture +def sigma_cutoff(): + """``sigma_cutoff(value)`` sets ``torchref.sigma_cutoff_ed`` and restores after.""" + original = torchref.sigma_cutoff_ed.value + + def _set(value: float) -> None: + torchref.sigma_cutoff_ed.value = value + + try: + yield _set + finally: + torchref.sigma_cutoff_ed.value = original + + +# --------------------------------------------------------------------------- +# Scenes +# --------------------------------------------------------------------------- +@pytest.fixture(scope="package") +def scene_small() -> H.Scene: + """5 atoms, few reflections -- sized for ``gradcheck``'s O(n_params) cost.""" + return H.synthetic_scene(n_atoms=5, a=14.0, d_min=3.0, max_refl=24, seed=3) + + +@pytest.fixture(scope="package") +def scene_fine() -> H.Scene: + """Synthetic P1 scene for parametrized coverage sweeps. + + 60 atoms rather than 10. Aliasing contamination in the *derivatives* cancels across + atoms roughly as 1/sqrt(N), so a 10-atom scene overstates it badly -- the deviatoric-U + gradnorm ratio is 1.43-1.70 at 10 atoms against 0.997 at 300 and 0.9995 on 7L84. See + the note on the gate constants in ``__init__.py``. + + Even at 60 atoms this remains pessimistic relative to a real structure, so it backs the + *comparative* tests -- backend parity, dtype coverage, cell angle -- while the absolute + accuracy gates are calibrated on ``gemmi_aniso_p1``. + """ + return H.synthetic_scene(n_atoms=60, a=24.0, d_min=1.6, max_refl=600, seed=0) + + +@pytest.fixture(scope="package") +def scene_monoclinic() -> H.Scene: + """beta = 115 deg. The non-orthogonal metric is where a node-centred or diagonal + truncation diverges most from the true Cartesian sphere.""" + return H.synthetic_scene( + n_atoms=60, a=24.0, beta=115.0, d_min=1.6, max_refl=600, seed=1 + ) + + +@pytest.fixture(scope="package") +def scene_coarse() -> H.Scene: + """Deliberately under-sampled: pins the sampling-dominated regime so the fine-grid + gate cannot pass by accident.""" + return H.synthetic_scene(n_atoms=60, a=24.0, d_min=3.5, max_refl=200, seed=0) + + +# --------------------------------------------------------------------------- +# gemmi scenes, from real deposited structures +# --------------------------------------------------------------------------- +_PDB_DIR = Path(__file__).resolve().parents[2] / "files" / "pdb" + + +@pytest.fixture(scope="package") +def gemmi_iso_p1(): + """3GR5 forced to P1. 1329 atoms after hydrogen removal.""" + pytest.importorskip("gemmi") + return H.gemmi_scene(_PDB_DIR / "3GR5.pdb", p1=True, d_min=3.0, max_refl=200) + + +@pytest.fixture(scope="package") +def gemmi_aniso_p1(): + """7L84 forced to P1 -- every one of its 1209 atoms carries ANISOU.""" + pytest.importorskip("gemmi") + return H.gemmi_scene(_PDB_DIR / "7L84.pdb", p1=True, d_min=3.0, max_refl=200) + + +@pytest.fixture(scope="package") +def gemmi_aniso_grad(): + """7L84 P1 at 1.5 A with enough reflections to calibrate derivative gates. + + Separate from ``gemmi_aniso_p1`` (which is sized for the cheap forward gemmi + comparison) because the absolute gradient and HVP gates must be set on a scene that is + representative, and representativeness here is driven by atom count and resolution. + """ + pytest.importorskip("gemmi") + return H.gemmi_scene(_PDB_DIR / "7L84.pdb", p1=True, d_min=1.5, max_refl=500) + + +@pytest.fixture(scope="package") +def oracle_aniso_grad(gemmi_aniso_grad): + """DS oracle for the real-structure scene: F, obs, gradients and an HVP.""" + scene, _ = gemmi_aniso_grad + return _oracle_bundle(scene) + + +@pytest.fixture(scope="package") +def gemmi_iso_symmetry(): + """3GR5 in its deposited P 65 2 2. + + Hexagonal on purpose. ``h' = h.R`` and ``h' = R.h`` coincide whenever the rotation + matrix is symmetric, so orthorhombic and tetragonal groups cannot distinguish them; + trigonal/hexagonal groups can. This is the same reasoning that drives the group + choice in ``tests/unit/symmetry/test_hkl_symmetry_gemmi.py``. + """ + pytest.importorskip("gemmi") + return H.gemmi_scene(_PDB_DIR / "3GR5.pdb", p1=False, d_min=3.0, max_refl=200) + + +# --------------------------------------------------------------------------- +# Memoized oracle results +# --------------------------------------------------------------------------- +@pytest.fixture(scope="package") +def oracle_fine(scene_fine): + """Forward ``F``, first-order grads and an HVP from the DS oracle, computed once.""" + return _oracle_bundle(scene_fine) + + +@pytest.fixture(scope="package") +def oracle_monoclinic(scene_monoclinic): + return _oracle_bundle(scene_monoclinic) + + +def _oracle_bundle(scene: H.Scene) -> dict: + """Forward ``F``, pseudo-observations, gradients and an HVP, all from the oracle. + + ``obs`` is derived here and handed to the tests, so the candidate and the oracle are + scored against *identical* pseudo-observations. Regenerating it per test would make + the two sides differentiate slightly different targets and the comparison would + measure that instead. + """ + out: dict = {} + v = _direction(scene) + for aniso in (False, True): + key = "aniso" if aniso else "iso" + fn = H.ds_aniso_oracle if aniso else H.ds_iso_oracle + + with torch.no_grad(): + F = fn(scene) + obs = H.synthetic_obs(F) + out[f"{key}_F"] = F + out[f"{key}_obs"] = obs + out[f"{key}_v"] = v + + xyz, occ, third = scene.leaves(aniso=aniso) + grads = torch.autograd.grad( + H.ls_target(fn(scene, xyz, occ, third), obs), (xyz, occ, third) + ) + out[f"{key}_grads"] = tuple(g.detach() for g in grads) + + xyz, occ, third = scene.leaves(aniso=aniso) + (g1,) = torch.autograd.grad( + H.ls_target(fn(scene, xyz, occ, third), obs), xyz, create_graph=True + ) + (hv,) = torch.autograd.grad((g1 * v).sum(), xyz) + out[f"{key}_hvp"] = hv.detach() + return out + + +def _direction(scene: H.Scene) -> torch.Tensor: + """Fixed unit direction for HVPs. Seeded, so the oracle and the candidate agree.""" + g = torch.Generator().manual_seed(17) + v = torch.randn(scene.xyz.shape, generator=g, dtype=scene.xyz.dtype) + return v / v.norm() diff --git a/tests/unit/structure_factor/helpers.py b/tests/unit/structure_factor/helpers.py new file mode 100644 index 00000000..c3856253 --- /dev/null +++ b/tests/unit/structure_factor/helpers.py @@ -0,0 +1,814 @@ +"""Shared scene construction, oracle wrappers and comparison metrics. + +Two scene sources: + +* :func:`synthetic_scene` -- a small random P1 scene, for the ``gradcheck`` / + ``gradgradcheck`` links where the cost is O(n_params) forward evaluations and the + scene size is what makes the difference between milliseconds and minutes. +* :func:`gemmi_scene` -- a real deposited structure, with the torchref scene built + **from** the gemmi structure so both sides are driven by the same atoms. + +That second point is load-bearing. If the two sides are populated in parallel from the +same file by two different readers, the test silently measures reader agreement instead +of structure-factor agreement, and a single differing atom shows up as a tolerance +problem rather than as the bug it is. So there is exactly one traversal of the gemmi +model, and torchref's tensors are filled from it. +""" + +from __future__ import annotations + +import contextlib + +import itertools +from dataclasses import dataclass, fields +from typing import Optional, Sequence + +import torch + +from torchref.utils import use_portable + +from torchref.base.direct_summation._backends import DS_BACKENDS +from torchref.base.direct_summation.dispatch import _eager_aniso, _eager_iso +from torchref.base.electron_density._backends import DENSITY_BACKENDS +from torchref.base.reciprocal import get_scattering_vectors, reciprocal_basis_matrix +from torchref.base.scattering.scattering_table import get_scattering_params_by_z +from torchref.symmetry.cell import Cell + +__all__ = [ + "Scene", + "GRID_FINENESS", + "synthetic_scene", + "gemmi_scene", + "gemmi_sf", + "ds_iso_oracle", + "ds_aniso_oracle", + "sf_fft_for", + "fft_sf", + "splat_kernel", + "device_supports_dtype", + "kernels_for", + "splat_direct", + "ds_kernel", + "ds_kernels_for", + "ds_direct", + "density_to_F", + "synthetic_obs", + "ls_target", + "best_fit_scale", + "rel_l2", + "max_rel", +] + +# Grids are built at ``max_res = d_min / GRID_FINENESS``, i.e. a voxel spacing of +# ``d_min / (3 * GRID_FINENESS)`` given ``NYQUIST_OVERSAMPLING = 3.0``. +# +# ``GRID_FINENESS = 1.0`` is the **production** configuration: ``max_res = d_min``, voxel +# spacing ``d_min / 3``. These tests default to it, because a gate on a finer grid than +# production ships would not bound anything that matters. +# +# Measured on ``scene_fine`` (60 atoms, a = 24 A, d_min 1.6), float64, against the DS +# oracle. Amplitudes are rel L2 on complex F; derivatives are of ``ls_target``: +# +# fineness spacing gridsize amplitude g_xyz g_xyz cos HVP HVP cos +# 0.667 d_min/2 (30, 36, 27) 1.04e-01 1.25e+00 0.4207 2.77e-01 0.9622 +# 1.000 d_min/3 (45, 50, 40) 4.11e-03 4.30e-02 0.999077 2.06e-02 0.999814 +# 1.300 d_min/3.9 (60, 64, 54) 8.01e-04 5.86e-03 0.999983 1.16e-03 0.999999 +# 1.600 d_min/4.8 (72, 80, 64) 8.03e-04 5.87e-03 0.999983 1.16e-03 0.999999 +# 2.200 d_min/6.6 (100,108, 90) 7.98e-04 6.00e-03 0.999982 1.16e-03 0.999999 +# +# Three things follow. +# +# 1. **Bare Nyquist is unusable**, which is why ``NYQUIST_OVERSAMPLING`` is 3 and not 2. At +# oversampling 2 the xyz gradient cosine against the analytic answer collapses to 0.42 +# and amplitudes are 10% out. The factor of 3 is buying a great deal. +# 2. **Production sits one step before convergence.** Everything is converged from fineness +# 1.3. At production the residuals are ~5x larger in amplitude and ~7x in the xyz +# gradient, but direction stays excellent (cos 0.999) so the residual is predominantly +# magnitude. On a *real* structure the production numbers are better still -- 7L84 gives +# amplitude 2.28e-03 and xyz gradient 1.04e-02 -- because derivative aliasing cancels +# across atoms as ~1/sqrt(N). See the gate constants in ``__init__.py``; absolute +# accuracy gates are calibrated there rather than here. +# 3. **The sigma cutoff is not the binding constraint at production.** Sweeping n_sigma at +# fineness 1.0 moves the amplitude residual 5.43e-3 -> 4.11e-3 -> 4.05e-3 and then +# flatlines; grid sampling dominates. Tests that mean to exercise the cutoff therefore +# pass an explicit finer ``fineness`` -- see +# ``test_forward.py::test_nsigma_reduces_truncation_error``. +GRID_FINENESS = 1.0 + + +# --------------------------------------------------------------------------- +# Scene +# --------------------------------------------------------------------------- +@dataclass(frozen=True) +class Scene: + """Everything both routes need, in one immutable bundle. + + ``xyz`` is Cartesian (what ``SfFFT`` / ``SfDS`` take) and ``xyz_frac`` is + fractional (what the ``direct_summation`` functions take). Both are carried rather + than converted at the call site, because getting that conversion wrong in one place + produces a plausible-looking wrong answer instead of an error. + """ + + cell: Cell + spacegroup: str + d_min: float + hkl: torch.Tensor # (R, 3) float + hkl_list: tuple # (R,) tuples of int, for gemmi + s: torch.Tensor # (R,) |s| = 1/d + s_vec: torch.Tensor # (R, 3) + xyz: torch.Tensor # (N, 3) Cartesian + xyz_frac: torch.Tensor # (N, 3) fractional + occ: torch.Tensor # (N,) + adp: torch.Tensor # (N,) B convention, A^2 + u6: torch.Tensor # (N, 6) [U11,U22,U33,U12,U13,U23], A^2 + A: torch.Tensor # (N, 5) ITC92 + B: torch.Tensor # (N, 5) ITC92 + frac_matrix: torch.Tensor # (3, 3) fractional -> Cartesian + inv_frac_matrix: torch.Tensor # (3, 3) Cartesian -> fractional + + @property + def n_atoms(self) -> int: + return int(self.xyz.shape[0]) + + @property + def n_refl(self) -> int: + return int(self.hkl.shape[0]) + + def leaves(self, *, aniso: bool = False, requires_grad: bool = True): + """Fresh differentiable copies of ``(xyz, occ, adp_or_u6)``.""" + third = self.u6 if aniso else self.adp + return tuple( + t.clone().requires_grad_(requires_grad) for t in (self.xyz, self.occ, third) + ) + + def to(self, device=None, dtype=None) -> "Scene": + """A new :class:`Scene` with every tensor moved and cast. + + Scenes are built CPU/float64 and stay that way, because the oracle is computed + from them and a reference must not inherit the precision of the thing it judges. + Only the *candidate* side is moved, so an MPS float32 kernel and a CPU float64 + oracle describe the same physical structure. + + ``hkl_list`` is untouched -- it is Python ints for gemmi, not a tensor. + """ + if device is None and dtype is None: + return self + real = {"hkl", "s", "s_vec", "xyz", "xyz_frac", "occ", "adp", "u6", "A", "B", + "frac_matrix", "inv_frac_matrix"} + moved = {} + for f in fields(self): + v = getattr(self, f.name) + moved[f.name] = ( + v.to(device=device, dtype=dtype) if f.name in real else v + ) + return Scene(**moved) + + +def _hkl_within(cell: Cell, d_min: float, dtype: torch.dtype, cap: Optional[int]): + """Every integer hkl inside the ``d_min`` shell, ordered, ``F(000)`` excluded. + + ``F(000)`` is dropped because the FFT route cannot reproduce it: the grid carries a + truncated density, so the zeroth term is the truncated electron count rather than + the true one. + """ + a, b, c = (float(cell.data[i]) for i in range(3)) + lim = [int(v / d_min) + 1 for v in (a, b, c)] + recB = reciprocal_basis_matrix(cell.data) + cand = [ + h + for h in itertools.product( + range(-lim[0], lim[0] + 1), + range(-lim[1], lim[1] + 1), + range(-lim[2], lim[2] + 1), + ) + if any(h) + ] + hkl = torch.tensor(cand, dtype=dtype) + s = get_scattering_vectors(hkl, cell.data, recB).norm(dim=1) + keep = (s > 0) & (s <= 1.0 / d_min) + hkl = hkl[keep] + if cap is not None and hkl.shape[0] > cap: + # Even stride, so the kept set still spans the full resolution range rather + # than clustering at low angle where agreement is easiest. + step = hkl.shape[0] // cap + 1 + hkl = hkl[::step] + return hkl + + +def _finish( + cell: Cell, + spacegroup: str, + d_min: float, + hkl: torch.Tensor, + z: torch.Tensor, + xyz_frac: torch.Tensor, + occ: torch.Tensor, + adp: torch.Tensor, + u6: torch.Tensor, + dtype: torch.dtype, + device: torch.device, +) -> Scene: + recB = reciprocal_basis_matrix(cell.data) + s_vec = get_scattering_vectors(hkl, cell.data, recB) + frac_matrix = cell.fractional_matrix.to(dtype) + inv_frac_matrix = cell.inv_fractional_matrix.to(dtype) + A, B = get_scattering_params_by_z(z, dtype=dtype) + to = lambda t: t.to(device=device, dtype=dtype) # noqa: E731 + return Scene( + cell=cell, + spacegroup=spacegroup, + d_min=d_min, + hkl=to(hkl), + hkl_list=tuple(tuple(int(v) for v in row) for row in hkl), + s=to(s_vec.norm(dim=1)), + s_vec=to(s_vec), + xyz=to(xyz_frac @ frac_matrix.T), + xyz_frac=to(xyz_frac), + occ=to(occ), + adp=to(adp), + u6=to(u6), + A=to(A), + B=to(B), + frac_matrix=to(frac_matrix), + inv_frac_matrix=to(inv_frac_matrix), + ) + + +def synthetic_scene( + *, + n_atoms: int = 10, + a: float = 24.0, + beta: float = 90.0, + d_min: float = 1.6, + seed: int = 0, + max_refl: Optional[int] = None, + dtype: torch.dtype = torch.float64, + device: torch.device = torch.device("cpu"), +) -> Scene: + """A small random P1 scene. + + Two degeneracies are deliberately avoided, both of which once survived in two test + files at the same time: + + * ``occ`` is never exactly 1.0. The kernels recover ``d/d_occ`` by dividing the + accumulated gradient by ``occ``, and at ``occ == 1`` that division is a no-op that + hides a wrong scaling. + * the ADP off-diagonals are non-zero **and signed**. Zero off-diagonals mean every + ellipsoid is axis-aligned, which leaves the cross-term arithmetic completely + uncovered -- the ``p01``/``p02``/``p12`` entries of the inverted 3x3, and the + backward's off-diagonal U gradients, which carry a ``4*pi^2`` factor where the + diagonal ones carry ``2*pi^2``. Magnitudes stay well below the diagonal so every U + remains comfortably positive-definite, since the Metal shader inverts ``M_g`` + analytically with no positive-definiteness guard. + + (Recorded here rather than in ``tests/helpers/kernel_cases.py``, which documented them + for the two accelerator test files this package replaced and has been removed with + them. ``test_dispatch.py::test_aniso_scene_exercises_off_diagonal_u`` asserts the + second one rather than trusting this docstring.) + """ + g = torch.Generator().manual_seed(seed) + cell = Cell( + torch.tensor([a, a * 1.1, a * 0.9, 90.0, beta, 90.0], dtype=torch.float64) + ) + hkl = _hkl_within(cell, d_min, torch.float64, max_refl) + + z = torch.tensor(([6, 7, 8, 16, 6, 8, 7, 26] * ((n_atoms // 8) + 1))[:n_atoms]) + # Keep atoms off the cell walls so a wrapped-vs-unwrapped bug in the map route + # cannot be confused with a truncation residual. + xyz_frac = 0.15 + 0.7 * torch.rand(n_atoms, 3, generator=g, dtype=torch.float64) + occ = 0.55 + 0.4 * torch.rand(n_atoms, generator=g, dtype=torch.float64) + adp = 8.0 + 22.0 * torch.rand(n_atoms, generator=g, dtype=torch.float64) + + u6 = torch.zeros(n_atoms, 6, dtype=torch.float64) + u6[:, :3] = (adp[:, None] / (8.0 * torch.pi**2)) * ( + 0.75 + 0.5 * torch.rand(n_atoms, 3, generator=g, dtype=torch.float64) + ) + off = 0.35 * u6[:, :3].min(dim=1).values + u6[:, 3:] = off[:, None] * ( + 2.0 * torch.rand(n_atoms, 3, generator=g, dtype=torch.float64) - 1.0 + ) + return _finish( + cell, "P 1", d_min, hkl, z, xyz_frac, occ, adp, u6, dtype, device + ) + + +def gemmi_scene( + pdb_path, + *, + p1: bool = True, + d_min: float = 3.0, + max_refl: Optional[int] = 200, + dtype: torch.dtype = torch.float64, + device: torch.device = torch.device("cpu"), +): + """Build a :class:`Scene` **from** a gemmi structure; return ``(scene, structure)``. + + Hydrogens are removed on the gemmi side before the traversal, so both routes see + the identical atom list -- ``remove_hydrogens`` after extraction would leave + torchref summing over atoms gemmi no longer has. + + ``p1=True`` rewrites the spacegroup to ``P 1`` and rebuilds ``cell.images``, which + is the mechanism gemmi's calculator uses for symmetry: ``calculate_sf_from_model`` + sums each atom over ``cell.images``, so an empty image list is a true P1 sum. + Verified: forcing P1 takes ``len(cell.images)`` from 11 to 0 on a P 65 2 2 entry. + """ + import gemmi + + st = gemmi.read_structure(str(pdb_path)) + st.setup_entities() + st.remove_hydrogens() + if p1: + st.spacegroup_hm = "P 1" + st.setup_cell_images() + + zs, frac, occ, adp, u6 = [], [], [], [], [] + for chain in st[0]: + for res in chain: + for at in res: + f = st.cell.fractionalize(at.pos) + an = at.aniso + zs.append(at.element.atomic_number) + frac.append([f.x, f.y, f.z]) + occ.append(at.occ) + adp.append(at.b_iso) + u6.append([an.u11, an.u22, an.u33, an.u12, an.u13, an.u23]) + + c = st.cell + cell = Cell( + torch.tensor([c.a, c.b, c.c, c.alpha, c.beta, c.gamma], dtype=torch.float64) + ) + hkl = _hkl_within(cell, d_min, torch.float64, max_refl) + scene = _finish( + cell, + st.spacegroup_hm, + d_min, + hkl, + torch.tensor(zs), + torch.tensor(frac, dtype=torch.float64), + torch.tensor(occ, dtype=torch.float64), + torch.tensor(adp, dtype=torch.float64), + torch.tensor(u6, dtype=torch.float64), + dtype, + device, + ) + return scene, st + + +def gemmi_sf(structure, hkl_list: Sequence, *, dtype=torch.complex128) -> torch.Tensor: + """``F(hkl)`` from gemmi's X-ray structure-factor calculator. + + ``addends`` is left untouched (zero), so there is no anomalous term to account for + on the torchref side. + """ + import gemmi + + calc = gemmi.StructureFactorCalculatorX(structure.cell) + return torch.tensor( + [complex(calc.calculate_sf_from_model(structure[0], h)) for h in hkl_list], + dtype=dtype, + ) + + +# --------------------------------------------------------------------------- +# Oracles +# --------------------------------------------------------------------------- +def ds_iso_oracle(scene: Scene, xyz=None, occ=None, adp=None) -> torch.Tensor: + """Isotropic ``F(hkl)`` from the pure-torch eager path. + + ``_eager_iso`` and not ``ds_iso``: the public entry points wrap ``_CheckpointedSF``, + whose backward omits ``create_graph=True``, so they raise on double backward. See + the package docstring. + + ``xyz`` is **Cartesian** here, matching :attr:`Scene.xyz`; it is fractionalized + inside so a differentiable leaf stays a leaf in Cartesian space, which is what the + FFT route differentiates with respect to. + """ + xyz = scene.xyz if xyz is None else xyz + occ = scene.occ if occ is None else occ + adp = scene.adp if adp is None else adp + return _eager_iso( + scene.hkl, + scene.s, + xyz @ scene.inv_frac_matrix.T, + occ, + adp, + scene.A, + scene.B, + None, + ) + + +def ds_aniso_oracle(scene: Scene, xyz=None, occ=None, u6=None) -> torch.Tensor: + """Anisotropic ``F(hkl)`` from the pure-torch eager path.""" + xyz = scene.xyz if xyz is None else xyz + occ = scene.occ if occ is None else occ + u6 = scene.u6 if u6 is None else u6 + return _eager_aniso( + scene.hkl, + scene.s_vec, + xyz @ scene.inv_frac_matrix.T, + occ, + u6, + scene.A, + scene.B, + None, + ) + + +def sf_fft_for( + scene: Scene, + dtype: torch.dtype = torch.float64, + *, + fineness: float = GRID_FINENESS, + spacegroup: str = "P 1", +): + """An ``SfFFT`` with its grid already set up, at the fineness policy above. + + Pass ``fineness=1.0`` for a deliberately under-sampled grid; see + :data:`GRID_FINENESS` for why that is the sampling-limited regime. + """ + from torchref.model.sf_fft import SfFFT + + sf = SfFFT( + cell=scene.cell, + spacegroup=spacegroup, + max_res=scene.d_min / fineness, + dtype_float=dtype, + device=torch.device("cpu"), + ) + sf.setup_grid() + return sf + + +# --------------------------------------------------------------------------- +# Direct kernel access +# --------------------------------------------------------------------------- +# Every production splat is called *directly* here rather than through +# ``build_electron_density``. Two reasons: +# +# 1. **No vacuity risk.** Under the default a failed accelerator kernel silently +# falls back to the portable splat (``main.py`` catches and falls through), so a +# dispatch-driven test can pass while measuring a different kernel than the one it +# names. Calling the kernel directly settles that by construction. +# 2. **No global-config coupling.** ``SfFFT`` builds its grid through ``get_real_grid``, +# which reads the *global* ``dtypes.float`` and takes no dtype argument -- so an MPS +# ``SfFFT`` under this package's float64 pin would try to allocate float64 on MPS and +# fail. ``ifft`` and ``extract_structure_factor_from_grid`` read no global config at +# all, so :func:`density_to_F` needs no config switching. +# +# The dispatch ladder is a separate concern, tested in ``test_dispatch.py``. + +#: The production density table, used directly rather than mirrored. +#: +#: This was a hand-maintained copy -- ``name, device, dtypes`` for each of the four splat +#: wrappers -- and copies of a policy drift. It already had: it encoded the six-tensor dtype +#: rule while the production Triton gate probed one tensor. Deriving the parametrization from +#: the table that ships means the matrix cannot claim coverage the dispatch does not have. +#: +#: All four wrappers share one signature, +#: ``(density_map, xyz, adp_or_u, occ, A, B, inv_frac, frac, radius_per_atom)``, so +#: :func:`splat_direct` needs no per-kernel adapter. They did not before the +#: standardization -- the Metal pair took extra unused arguments in a different order, +#: and CUDA had no wrapper at all. +_KERNEL_TABLE = DENSITY_BACKENDS + + +def splat_kernel(name: str, aniso: bool): + """Resolve a kernel name to its wrapper, via the production table. + + ``Backend.resolve`` does the import lazily, which is what a CPU-only host needs at + collection time: the Metal module loads MSL source and the CUDA one imports Triton. + """ + return _KERNEL_TABLE.by_name(name).resolve(aniso) + + +def maybe_portable(pin: bool): + """``use_portable()`` when ``pin``, otherwise a no-op context. + + Lets a test parametrize over ``[False, True]`` and get the default backend versus the + pinned reference from one ``with``. Replaces the ``use_engine(engine)`` idiom that the + ``[Engine.AUTO, Engine.EAGER]`` parametrizations used. + + Those parametrizations are all still two *distinct* kernels, which is worth stating + because it is easy to assume otherwise: on CPU the default selects the fused C++ sphere + splat and the pin selects the portable ``scatter_add`` one. Even at second order they + differ -- only the double-backward *re-derivation* borrows the portable splat, while the + forward and first derivative stay in C++. + """ + return use_portable() if pin else contextlib.nullcontext() + + +def device_supports_dtype(device: torch.device, dtype: torch.dtype) -> bool: + """Whether ``device`` can hold ``dtype`` at all. + + MPS has no float64 — it is a *device* limitation, not a kernel one, so it belongs + here and not in the per-kernel table. The portable splat is float64-capable and + device-agnostic, and without this check it would be offered for ``(mps, float64)`` + and fail inside ``Scene.to``. + """ + if device.type == "mps" and dtype is torch.float64: + return False + return True + + +def _table_names_for(table, device: torch.device, dtype: torch.dtype): + """Backends in ``table`` whose declared device/dtype envelope covers this pair. + + Reads ``Backend.device`` / ``Backend.dtypes`` rather than calling ``select``: the point + is to enumerate every kernel that *can* run here so each gets its own direct-call + accuracy test, whereas ``select`` answers the narrower question of which one wins. + """ + if not device_supports_dtype(device, dtype): + return [] + out = [] + for b in table.backends: + if b.device is not None and b.device != device.type: + continue + if b.dtypes is not None and dtype not in b.dtypes: + continue + out.append(b.name) + return out + + +def kernels_for(device: torch.device, dtype: torch.dtype): + """Kernel names that actually run on this ``(device, dtype)``. + + Filtering here rather than skipping inside the test keeps a wrong entry visible: an + unsupported combination produces no test rather than a passing one. + """ + return _table_names_for(_KERNEL_TABLE, device, dtype) + + +def splat_direct(scene: Scene, name: str, xyz=None, occ=None, third=None, *, aniso=False): + """Call one splat kernel directly and return the density map. + + ``scene`` must already be on the target device and dtype (see :meth:`Scene.to`). + The per-atom radius is computed here because the kernels take it as an argument -- + ``build_electron_density`` normally does this, and bypassing the dispatch means the + caller owns it. Same policy call the dispatch makes, so the truncation contract is + unchanged. + """ + from torchref.base.electron_density.radius_policy import ( + per_atom_radius_aniso, + per_atom_radius_iso, + ) + from torchref.config import get_sigma_cutoff_ed + + xyz = scene.xyz if xyz is None else xyz + occ = scene.occ if occ is None else occ + if third is None: + third = scene.u6 if aniso else scene.adp + + n_sigma = get_sigma_cutoff_ed() + radius = ( + per_atom_radius_aniso(scene.B, third, n_sigma=n_sigma) + if aniso + else per_atom_radius_iso(third, scene.B, n_sigma=n_sigma) + ) + dims = _grid_dims(scene) + density_map = torch.zeros(*dims, dtype=xyz.dtype, device=xyz.device) + fn = splat_kernel(name, aniso) + return fn( + density_map, xyz, third, occ, scene.A, scene.B, + scene.inv_frac_matrix, scene.frac_matrix, radius, + ) + + +# --------------------------------------------------------------------------- +# Direct-summation kernels, also called directly +# --------------------------------------------------------------------------- +#: The production direct-summation table, plus ``eager`` -- which is deliberately *not* in +#: it. ``_eager_*`` is the double-differentiable oracle, not a dispatch target, so it has no +#: business in a table describing what production selects; it is resolved here instead. +#: +#: There is **no Metal/MPS direct-summation kernel**, which the table now says outright by +#: there being no Metal DS kernel at all. So DS on MPS *is* +#: ``_checkpointed_*`` running on-device -- a real production path, and what the ``mps`` leg +#: of ``checkpointed`` covers. +_DS_KERNEL_TABLE = DS_BACKENDS + + +def ds_kernel(name: str, aniso: bool): + """Resolve a direct-summation kernel name to its function. + + ``Backend.resolve`` keeps the table's imports lazy: the Triton path pulls in ``triton``, + which a CPU-only host should not need at collection time. + """ + if name == "eager": + return _eager_aniso if aniso else _eager_iso + return _DS_KERNEL_TABLE.by_name(name).resolve(aniso) + + +def ds_kernels_for(device: torch.device, dtype: torch.dtype): + """Candidate DS kernels on this ``(device, dtype)``. + + Excludes the oracle by construction -- it is not in the production table. + """ + return _table_names_for(_DS_KERNEL_TABLE, device, dtype) + + +def ds_direct(scene: Scene, name: str, xyz_frac=None, occ=None, third=None, *, aniso=False): + """Call one direct-summation kernel directly and return complex ``F(hkl)``. + + Coordinates are **fractional** here -- unlike the splat kernels, which take Cartesian. + Carrying both on :class:`Scene` is what keeps that from being a silent error. + + ``ds_iso_triton`` takes no ``max_memory_gb``; the eager and checkpointed paths do. That + is the only signature difference, and it is absorbed here rather than at call sites. + """ + xyz_frac = scene.xyz_frac if xyz_frac is None else xyz_frac + occ = scene.occ if occ is None else occ + if third is None: + third = scene.u6 if aniso else scene.adp + geom = scene.s_vec if aniso else scene.s + fn = ds_kernel(name, aniso) + args = (scene.hkl, geom, xyz_frac, occ, third, scene.A, scene.B) + if name == "ds_triton": + return fn(*args) + return fn(*args, None) + + +def _grid_dims(scene: Scene, fineness: float = GRID_FINENESS): + """Grid dimensions matching what ``SfFFT.setup_grid`` would choose. + + Derived arithmetically instead of by constructing an ``SfFFT``, so no global dtype is + read and nothing float64 is allocated on a device that cannot hold it. Rounded to + even numbers, which is all the kernels care about (they take the shape from + ``density_map``); FFT-friendliness only matters for speed here. + """ + spacing = scene.d_min / (3.0 * fineness) + lengths = [float(scene.cell.data[i]) for i in range(3)] + return tuple(max(4, 2 * int(round(L / spacing / 2.0))) for L in lengths) + + +def density_to_F(scene: Scene, density_map: torch.Tensor) -> torch.Tensor: + """``F(hkl)`` from a density map, without ``SfFFT``. + + ``ifft`` applies the ``V_cell / N`` voxel-volume scaling that makes the result + directly comparable to a direct-summation atom sum -- no scale factor, no offset. + Neither this nor ``extract_structure_factor_from_grid`` reads the global dtype + config, which is what lets an accelerator candidate run while the package is pinned + to float64 for the oracle. + """ + from torchref.base.fourier.fft import ifft + from torchref.base.reciprocal.grid_operations import ( + extract_structure_factor_from_grid, + ) + + volume = float(scene.cell.volume) + return extract_structure_factor_from_grid( + ifft(density_map, volume), scene.hkl + ) + + +def fft_sf(scene: Scene, sf_fft, xyz=None, occ=None, third=None, *, aniso=False): + """``F(hkl)`` through the density-splat + FFT route, in P1. + + ``apply_symmetry=False`` on both calls: P1 isolates the truncation and sampling + budget, and a symmetric comparison would cancel the symmetry algebra anyway, since + both routes call the same ``compute_symmetry_equivalent_hkls`` / + ``compute_translation_phases``. Symmetry is validated against gemmi instead, in + ``test_forward.py``. + """ + xyz = scene.xyz if xyz is None else xyz + occ = scene.occ if occ is None else occ + if third is None: + third = scene.u6 if aniso else scene.adp + + # Match the grid's dtype. Scenes are built in float64 so the DS oracle keeps full + # precision, but production runs float32 and ``SfFFT`` derives its cell matrices from + # ``dtype_float`` -- feeding float64 atoms into a float32 grid raises inside + # ``_canonical_setup``. Casting is differentiable, so gradient tests still reach the + # original float64 leaf. + target = sf_fft.dtype_float + if xyz.dtype != target: + xyz, occ, third = xyz.to(target), occ.to(target), third.to(target) + A, B = scene.A.to(target), scene.B.to(target) + + empty1 = xyz.new_zeros(0) + empty3 = xyz.new_zeros(0, 3) + empty5 = A.new_zeros(0, 5) + if aniso: + dm = sf_fft.build_density_map( + empty3, + empty1, + empty1, + empty5, + empty5, + xyz_aniso=xyz, + u_aniso=third, + occ_aniso=occ, + A_aniso=A, + B_aniso=B, + apply_symmetry=False, + ) + else: + dm = sf_fft.build_density_map(xyz, third, occ, A, B, apply_symmetry=False) + return sf_fft.map_to_structure_factors(dm, scene.hkl, apply_symmetry=False) + + +# --------------------------------------------------------------------------- +# Metrics +# --------------------------------------------------------------------------- +def synthetic_obs( + F_ref: torch.Tensor, *, rel_sigma: float = 0.10, seed: int = 101 +) -> torch.Tensor: + """Pseudo-observed amplitudes: ``|F_ref|`` perturbed by ``rel_sigma`` relative noise. + + This exists to make :func:`ls_target` well-posed. A least-squares residual taken + directly against ``|F_ref|`` has its minimum exactly where the two routes agree, so + near convergence the gradient comparison becomes partly self-referential -- the + quantity being differentiated is the very difference under test. Offsetting the + target by a fixed random 10% breaks that: the residual is O(0.1|F|), dominated by + synthetic misfit rather than by the FFT-vs-DS discrepancy, which is the regime a + refinement actually operates in. + + Deterministic by seed, and generated once from the oracle, so the candidate and the + oracle are scored against the *same* target. Multiplicative rather than additive + noise, so weak reflections are not swamped. + """ + g = torch.Generator().manual_seed(seed) + amp = F_ref.abs() + noise = torch.randn(amp.shape, generator=g, dtype=amp.dtype) + return (amp * (1.0 + rel_sigma * noise)).clamp_min(0.0) + + +def ls_target( + F: torch.Tensor, f_obs: torch.Tensor, weights: Optional[torch.Tensor] = None +) -> torch.Tensor: + """Least-squares amplitude residual ``sum(w * (|F| - f_obs)^2)``. + + The quantitative target for the gradient and HVP gates. Unweighted by default, so it + is dominated by strong reflections in the way an unweighted crystallographic LS + residual is. + + Well-conditioned in a way :func:`phase_sensitive_loss` is not: the per-reflection + residuals are random-signed at O(0.1|F|), so the sum grows like sqrt(N) instead of + cancelling, and the relative-error metric measures accuracy rather than the + conditioning of a near-cancelled sum. + """ + r = F.abs().to(f_obs.dtype) - f_obs + return (r * r).sum() if weights is None else (weights * r * r).sum() + + +# Why there is no phase-sensitive loss here +# ----------------------------------------- +# An earlier draft used ``(F.real**2 + 2*F.imag).sum()`` -- asymmetric in the imaginary +# part -- on the theory that a squared target cannot detect a kernel that conjugates F. +# Both premises were wrong. +# +# It is not testable this way. A crystallographic amplitude target is phase-blind by +# construction, and so are its derivatives: +# +# d|F|/dx = Re(F* dF/dx) / |F| +# +# Conjugate the whole calculation -- F -> F*, dF/dx -> (dF/dx)* -- and +# Re(F* dF/dx) -> Re(F (dF/dx)*) = Re((F* dF/dx)*) = Re(F* dF/dx), unchanged. A global +# phase convention is an unobservable gauge choice for such a target, so asserting a +# particular imaginary sign through it tests nothing physical. +# +# And it actively harmed the measurement. A linear term in F.imag sums signed quantities +# that largely cancel, so its gradient is a small difference of large numbers and every +# per-reflection error is amplified in relative terms. Measured on the same FFT-vs-DS xyz +# gradient comparison: 2.7e-01 under the asymmetric functional against ~6e-03 under +# ``ls_target`` -- a 40x artefact of conditioning, not of accuracy. +# +# Phase *convention* still matters, because anomalous scattering and difference maps are +# phase-sensitive and because the two routes must agree with each other. It is gated where +# it is observable: on the forward complex F, against gemmi and between routes, in +# ``test_forward.py``. See ``test_ls_target_is_phase_blind`` there, which pins the +# invariance above so this reasoning cannot quietly rot. + + +def rel_l2(got: torch.Tensor, ref: torch.Tensor) -> float: + return float((got - ref).norm() / ref.norm()) + + +def max_rel(got: torch.Tensor, ref: torch.Tensor, *, significance: float = 1e-3) -> float: + """Worst per-reflection relative error, over *significant* reflections only. + + Reflections with ``|F_ref| <= significance * max|F_ref|`` are excluded. Without that + filter this metric is dominated by systematic absences: in ``P 65 2 2`` the + extinguished reflections have ``|F|`` at round-off, so a ratio there reports ~6e-3 + while the rel-L2 over the same set is 4e-8. That is a property of dividing by zero, + not a disagreement, and gating on it would mean either a meaningless tolerance or + deleting the symmetry test. + + Reported alongside :func:`rel_l2` rather than instead of it -- rel L2 can hide a + single badly wrong strong reflection, and this cannot. + """ + scale = ref.abs().max() + keep = ref.abs() > significance * scale + if not bool(keep.any()): + return 0.0 + return float(((got[keep] - ref[keep]).abs() / ref[keep].abs()).max()) + + +def best_fit_scale(got: torch.Tensor, ref: torch.Tensor) -> float: + """Least-squares scale ``k`` minimizing ``||k*|ref| - |got|||``. + + Compared on amplitudes so a phase convention difference does not leak into the + scale estimate; phase is gated separately. + """ + g, r = got.abs().reshape(-1), ref.abs().reshape(-1) + return float((r * g).sum() / (r * r).sum()) diff --git a/tests/unit/structure_factor/test_dispatch.py b/tests/unit/structure_factor/test_dispatch.py new file mode 100644 index 00000000..7283154f --- /dev/null +++ b/tests/unit/structure_factor/test_dispatch.py @@ -0,0 +1,384 @@ +"""Dispatch end to end, as distinct from the selection *decision*. + +Deliberately separate from the accuracy tests. Those call each kernel **directly**, so they +say nothing about whether ``build_electron_density`` would have chosen it — and that +separation is the point. + +Also distinct from ``tests/unit/utils/test_backend_tables.py``, which asserts what ``select`` +returns. That is a pure function and can be interrogated for any ``(device, dtype)``; what it +cannot show is that the chosen kernel was then actually *called*. So what is pinned here is +provenance and fallback behaviour, both of which need a real dispatch on a real device: + +* on an accelerator host, the default really reaches the native kernel rather than the + portable splat (call recorder, patched at the kernel's defining module); +* ``force_portable`` reaches the portable splat on every device; +* with an accelerator forced dead, the fallback produces a numerically correct map. + +Ported from ``tests/integration/test_variable_radius_{gpu,mps}.py``, which are deleted: their +accuracy coverage is superseded by the oracle legs in this package, but these dispatch +contracts are not accuracy and had no replacement. +""" + +from __future__ import annotations + +import pytest +import torch + +from torchref.base.electron_density.main import build_electron_density +from torchref.base.electron_density.radius_policy import per_atom_radius_iso + +from tests.helpers.grad_asserts import rel_error + +from . import RTOL_DS_F32 +from . import helpers as H + +pytestmark = pytest.mark.unit + + +def _build(scene, device, dtype, force_portable=False, aniso=False): + """Run one dispatch through ``build_electron_density`` on ``device``. + + ``dtype`` is passed explicitly rather than inherited from the ambient config. That is + load-bearing for the accelerator branches: under this package's float64 pin the Metal + and Triton gates would never fire, which is a silent way for every assertion below to + become vacuous. + """ + s = scene.to(device=device, dtype=dtype) + dims = H._grid_dims(s) + grid = torch.zeros(*dims, 3, dtype=dtype, device=device) + voxel = torch.tensor( + [float(s.cell.data[i]) / dims[i] for i in range(3)], dtype=dtype, device=device + ) + empty1 = s.xyz.new_zeros(0) + empty3 = s.xyz.new_zeros(0, 3) + empty5 = s.A.new_zeros(0, 5) + kw = dict( + real_space_grid=grid, + inv_frac_matrix=s.inv_frac_matrix, + frac_matrix=s.frac_matrix, + voxel_size=voxel, + dtype=dtype, + ) + kw["force_portable"] = force_portable + if aniso: + return build_electron_density( + xyz_iso=empty3, adp_iso=empty1, occ_iso=empty1, + A_iso=empty5, B_iso=empty5, + xyz_aniso=s.xyz, u_aniso=s.u6, occ_aniso=s.occ, + A_aniso=s.A, B_aniso=s.B, **kw, + ) + return build_electron_density( + xyz_iso=s.xyz, adp_iso=s.adp, occ_iso=s.occ, + A_iso=s.A, B_iso=s.B, **kw, + ) + + +# --------------------------------------------------------------------------- +# Provenance: the named kernel actually runs +# --------------------------------------------------------------------------- +@pytest.mark.mps +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +def test_metal_kernel_is_actually_dispatched(scene_small, kind, monkeypatch): + """On MPS float32, the default must call the Metal kernel. + + Verified with a call recorder, not by comparing maps: the portable reference uses + ``scatter_add`` on MPS, whose accumulation order is not reproducible, so two portable + runs already differ at ~1e-7 and an equality check against one would prove nothing + either way. + + **The patch target is the defining module**, and it is the same rule for every backend: + dispatch resolves each kernel by ``(module_path, attr)`` at call time, so the name to + patch is the one the table points at. That uniformity is new -- this test and the CUDA + one used to need different targets (a package here, ``main`` there) because the two + ladders bound their kernels differently. + + The recorder delegates to the real kernel rather than stubbing it, so the dispatch is + still exercised end to end. + """ + from torchref.base.electron_density.kernels.mps import variable_radius as kernels_mps + + name = f"add_{'anisotropic' if kind == 'aniso' else 'isotropic'}_mps_var" + real = getattr(kernels_mps, name) + calls = [] + + def recording(*args, **kwargs): + calls.append(1) + return real(*args, **kwargs) + + monkeypatch.setattr(kernels_mps, name, recording) + _build(scene_small, torch.device("mps"), torch.float32, aniso=kind == "aniso") + assert calls, f"the default did not dispatch to {name}" + + +@pytest.mark.cuda +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +def test_triton_kernel_is_actually_dispatched(scene_small, kind, monkeypatch): + """CUDA float32 equivalent of the Metal provenance test. + + Same patch target rule as the Metal case -- the module that defines the kernel. The two + used to differ, which is why they are separate tests rather than one parametrized + helper; that reason is gone, and merging them is a follow-up worth doing on a host that + has both backends to run it on. + """ + from torchref.base.electron_density.kernels.cuda import variable_radius as kernels_cuda + + name = f"add_{'anisotropic' if kind == 'aniso' else 'isotropic'}_cuda_var" + real = getattr(kernels_cuda, name) + calls = [] + + def recording(*args, **kwargs): + calls.append(1) + return real(*args, **kwargs) + + monkeypatch.setattr(kernels_cuda, name, recording) + _build(scene_small, torch.device("cuda"), torch.float32, aniso=kind == "aniso") + assert calls, f"the default did not dispatch to {name}" + + +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +def test_force_portable_reaches_the_portable_splat(scene_small, kind, monkeypatch): + """``force_portable`` must reach the portable splat, on any device. + + The counterpart to the provenance tests above, and the only override the dispatcher has: + it exists so you can pin the reference implementation when you suspect an accelerator is + returning wrong numbers, so it has to be the *portable* kernel that runs and not whatever + the default would have picked. + """ + from torchref.base.electron_density.kernels.cpu import ( + variable_radius as kernels_portable, + ) + + name = f"add_{'anisotropic' if kind == 'aniso' else 'isotropic'}_plain_var" + real = getattr(kernels_portable, name) + calls = [] + + def recording(*args, **kwargs): + calls.append(1) + return real(*args, **kwargs) + + monkeypatch.setattr(kernels_portable, name, recording) + _build(scene_small, torch.device("cpu"), torch.float32, force_portable=True, + aniso=kind == "aniso") + assert calls, f"force_portable did not dispatch to {name}" + + +# --------------------------------------------------------------------------- +# Failure policy: strict engines raise, AUTO degrades +# --------------------------------------------------------------------------- +@pytest.mark.mps +def test_default_degrades_correctly_when_the_shader_is_unavailable(scene_small, monkeypatch): + """With the shader dead, dispatch must still produce a *correct* map. + + Selection tests prove the *decision*; this proves the fallback's numbers. It is the only + place that runs the degradation end to end rather than asserting which row was chosen. + + The half of this test that asserted ``Engine.METAL`` raises is gone with the enum. Its + job -- making a dead accelerator loud rather than silently slow -- is now done by two + other things: ``test_backend_is_available_where_it_is_expected`` fails on any run on an + MPS host where the shader should work and doesn't, and the degradation warning is promoted + to an error under pytest. Both are strictly earlier and broader than a forced engine was. + + ``monkeypatch`` reverts the module globals at teardown, so this needs no + ``try``/``finally`` -- and must not have one, since a swallowed failure here is precisely + the behaviour under test. + """ + from torchref.base.electron_density.kernels.mps import compile as mps_compile + + monkeypatch.setattr(mps_compile, "_lib", None) + monkeypatch.setattr(mps_compile, "_lib_failed", True) + monkeypatch.setattr(mps_compile, "_lib_error", ("forced test failure", "")) + + mps = torch.device("mps") + # The default degrades to a working splat. Compared against the pinned portable path + # device, not against the oracle: with the shader forced to fail both engines now reach + # the *same* portable kernel, so they must agree closely, and an oracle comparison here + # would instead be measuring this scene's discretization error (which is large -- it is + # deliberately tiny and coarse for gradcheck) and tell us nothing about the fallback. + # Loose rather than bitwise because MPS ``scatter_add`` accumulation order is not + # reproducible, so two portable runs already differ at ~1e-7. + auto = _build(scene_small, mps, torch.float32) + portable = _build(scene_small, mps, torch.float32, force_portable=True) + rel = H.rel_l2(auto.cpu().to(torch.float64), portable.cpu().to(torch.float64)) + print(f"\n default vs portable after forced shader failure: relL2 {rel:.3e}") + assert rel < 1e-5, ( + f"dispatch did not degrade to the portable splat after the shader failed " + f"(rel {rel:.3e})" + ) + + +# --------------------------------------------------------------------------- +# Non-vacuity +# --------------------------------------------------------------------------- +def test_aniso_scene_exercises_off_diagonal_u(scene_small, scene_fine): + """The anisotropic scenes must have non-zero off-diagonal U. + + Ported from ``test_variable_radius_gpu.py``. Every aniso comparison in this package + would pass on axis-aligned ellipsoids while leaving the cross-term arithmetic + untouched -- the ``p01``/``p02``/``p12`` entries of the inverted 3x3, and the backward's + off-diagonal U gradients, which carry a ``4*pi^2`` factor where the diagonal ones carry + ``2*pi^2``. A scene-level assertion, since the scenes are what the tests share. + """ + for name, scene in (("scene_small", scene_small), ("scene_fine", scene_fine)): + off = scene.u6[:, 3:] + assert off.abs().max() > 0, f"{name}: off-diagonal U is all zero" + assert (off < 0).any() and (off > 0).any(), ( + f"{name}: off-diagonal U is single-signed; a sign error in the cross terms " + "would not show up" + ) + + +def test_radius_policy_is_the_same_on_every_device(scene_small): + """The per-atom truncation radius must not depend on device or dtype. + + ``splat_direct`` computes the radius itself because it bypasses the dispatch. If that + ever diverged from what ``build_electron_density`` computes, every accuracy test in + this package would be measuring a different truncation than production uses -- and + would still pass, because both sides would share it. + """ + from torchref.config import get_sigma_cutoff_ed + + ns = get_sigma_cutoff_ed() + ref = per_atom_radius_iso(scene_small.adp, scene_small.B, n_sigma=ns) + for device in (torch.device("cpu"), torch.device("mps")): + if device.type == "mps" and not torch.backends.mps.is_available(): + continue + s = scene_small.to(device=device, dtype=torch.float32) + got = per_atom_radius_iso(s.adp, s.B, n_sigma=ns) + assert torch.equal(got.cpu().to(ref.dtype), ref), ( + f"radius policy differs on {device.type}: the truncation contract is not " + "device-independent" + ) + + +# --------------------------------------------------------------------------- +# Under-probed gates (CUDA only) +# --------------------------------------------------------------------------- +# Both of these assert the *intended* contract, not current behaviour. They were written +# from reading the gates on a host with no CUDA, so they have never run. If one fails on a +# GPU host, that is the finding -- do not loosen it to match what the code does. + + +@pytest.mark.cuda +def test_triton_density_gate_probes_every_tensor(scene_small): + """A float64 ``density_map`` must not reach the float32 Triton kernel. + + Written when the density Triton branch gated on ``should_use_triton(xyz)`` -- **only + xyz** -- while its sibling gates probed six tensors. A float32 ``xyz`` with a float64 + ``density_map`` therefore passed, and ``WorkQueueGridDensity.forward`` handed that + float64 buffer to a kernel doing float32 ``tl.atomic_add``. + + The ``cuda_triton`` row now declares ``probes`` covering all six, so the mixed set + resolves to ``portable`` instead and this should never reach the Triton kernel at all. + Kept as an end-to-end guard rather than deleted: the fix is a table entry, and a table + entry can be edited. It has still never executed -- there is no CUDA host here -- so + treat a failure as information, not as a broken test. + + The contract asserted is that such a call fails loudly. Either outcome is acceptable -- + a refusal at selection, or a raise from the kernel -- but silently returning a number is + not. + """ + cuda = torch.device("cuda") + s = scene_small.to(device=cuda, dtype=torch.float32) + dims = H._grid_dims(s) + dm_f64 = torch.zeros(*dims, dtype=torch.float64, device=cuda) + from torchref.base.electron_density.kernels.cuda.variable_radius import ( + add_isotropic_cuda_var, + ) + from torchref.config import get_sigma_cutoff_ed + + radius = per_atom_radius_iso(s.adp, s.B, n_sigma=get_sigma_cutoff_ed()) + + with pytest.raises(Exception): + out = add_isotropic_cuda_var( + dm_f64, s.xyz, s.adp, s.occ, s.A, s.B, + s.inv_frac_matrix, s.frac_matrix, radius, + ) + # If it did not raise, it must at least not have silently produced garbage. + assert torch.isfinite(out).all() and out.abs().sum() > 0, ( + "float64 density_map + float32 Triton kernel returned a non-finite or empty " + "map instead of raising" + ) + + +@pytest.mark.cuda +def test_triton_ds_does_not_silently_truncate_hkl(scene_small): + """The Triton DS kernel must not silently downcast a float64 ``hkl``. + + ``triton_ds.py`` gates on ``xyz_frac`` alone and then force-casts every other input to + float32 (``_cols_f32``), including ``hkl``. ``hkl`` feeds the phase + ``phi = 2*pi*(h.r)``, the most precision-sensitive quantity in the calculation, and a + float64 ``hkl`` is exactly what a float64 config produces (``sf_ds.py:473``). + + Asserted against the eager oracle rather than against the float32 result: the question + is not whether the two Triton calls agree with each other, it is whether the truncation + costs accuracy that the caller asked for by supplying float64. + """ + from torchref.base.direct_summation.dispatch import _eager_iso + from torchref.base.direct_summation.triton_ds import ds_iso_triton + + cuda = torch.device("cuda") + s64 = scene_small.to(device=cuda, dtype=torch.float64) + s32 = scene_small.to(device=cuda, dtype=torch.float32) + + ref = _eager_iso( + s64.hkl, s64.s, s64.xyz_frac, s64.occ, s64.adp, s64.A, s64.B, None + ) + # float64 hkl with float32 xyz_frac: the gate passes, and hkl gets truncated. + got = ds_iso_triton( + s64.hkl, s32.s, s32.xyz_frac, s32.occ, s32.adp, s32.A, s32.B + ) + rel = H.rel_l2(got.cpu().to(torch.complex128), ref.cpu().to(torch.complex128)) + print(f"\n Triton DS with float64 hkl: relL2 vs eager oracle {rel:.3e}") + assert rel < 1e-4, ( + f"a float64 hkl into the Triton DS kernel costs {rel:.3e} against the oracle. The " + "kernel truncates it to float32 with no warning; either it should preserve the " + "precision or the gate should refuse the call." + ) + + +@pytest.mark.cuda +def test_sfds_backend_toggle_end_to_end(scene_fine): + """``SfDS`` through a full symmetry loop: the Triton DS kernel vs the reference. + + Restores the end-to-end coverage lost with ``test_ds_triton_vs_eager.py``. Distinct from + the per-kernel DS legs in ``test_forward.py`` / ``test_gradients.py``: those call the + kernels directly in P1, so nothing there exercises ``SfDS``'s symmetry accumulation or its + per-call ``force_portable`` argument. + + A non-P1 group on purpose. Both backends share the symmetry algebra, so this does not + validate the convention -- ``test_forward.py::test_sfds_matches_gemmi_with_symmetry`` does + that against gemmi. What it validates is that swapping the backend underneath a symmetry + loop changes nothing. + + Non-vacuity does not depend on being able to *force* Triton: on a CUDA host + ``test_backend_is_available_where_it_is_expected`` fails if ``ds_triton`` should work and + does not, and a runtime fallback raises via the degradation warning. Both fire before this + test could quietly compare the reference against itself. + """ + from torchref.model.sf_ds import SfDS + + cuda = torch.device("cuda") + s = scene_fine.to(device=cuda, dtype=torch.float32) + obs = H.synthetic_obs(H.ds_direct(scene_fine, "eager").detach()).to(cuda, torch.float32) + + def run(force_portable): + sf = SfDS( + cell=s.cell, spacegroup="P212121", force_portable=force_portable, + dtype_float=torch.float32, device=cuda, max_memory_gb=2.0, + ) + leaves = tuple(t.clone().requires_grad_(True) for t in (s.xyz, s.adp, s.occ)) + xyz, adp, occ = leaves + F, _ = sf.compute_structure_factors(s.hkl, xyz, adp, occ, s.A, s.B) + grads = torch.autograd.grad(H.ls_target(F, obs), leaves) + return F.detach(), grads + + F_t, g_t = run(False) # default: the Triton DS kernel on CUDA float32 + F_e, g_e = run(True) # pinned to the checkpointed reference + + rel_F = H.rel_l2(F_t.cpu().to(torch.complex128), F_e.cpu().to(torch.complex128)) + print(f"\n SfDS P212121 triton vs portable: F relL2 {rel_F:.3e}") + assert rel_F < RTOL_DS_F32, f"SfDS forward differs by engine: rel {rel_F:.3e}" + for name, a, b in zip(("xyz", "adp", "occ"), g_t, g_e): + rel = rel_error(a, b) + print(f" {name:4s} grad relL2 {rel:.3e}") + assert rel < RTOL_DS_F32, f"SfDS {name} gradient differs by engine: rel {rel:.3e}" diff --git a/tests/unit/structure_factor/test_ds_dispatch.py b/tests/unit/structure_factor/test_ds_dispatch.py new file mode 100644 index 00000000..e329403b --- /dev/null +++ b/tests/unit/structure_factor/test_ds_dispatch.py @@ -0,0 +1,154 @@ +"""Direct-summation dispatch *mechanics*, as distinct from its accuracy. + +CPU-only. What is left here after the oracle package absorbed the numerics: that +reflection-chunking is exact, that an empty atom set returns zeros of the right shape, and +that an integer ``hkl`` -- the production dtype -- is accepted losslessly. + +Two tests were deleted rather than migrated when the ``Engine`` enum went away. Both had +forcing as their *subject*: one asserted that a forced Triton engine refuses CPU inputs, and +one that a forced Metal engine runs eager instead of raising. With no forcing engine there is +nothing to refuse, and the second would have degenerated into ``assert torch.equal(F, F)``. +Selection is now asserted directly against the tables in +``tests/unit/utils/test_backend_tables.py``. + +The accuracy questions this file used to own -- checkpointed-vs-eager parity and +``gradcheck`` -- moved to ``test_gradients.py`` in this package, which ties the shipping +``_checkpointed_*`` path to the ``_eager_*`` oracle and then to finite differences. They +were duplicated in ``tests/unit/test_gradient_correctness.py`` as well, at a different +strictness, which is why they were consolidated. +""" + +import pytest +import torch + +from torchref.base.direct_summation import ds_iso +from torchref.base.direct_summation import dispatch as D +from torchref.config import dtypes + +pytestmark = pytest.mark.unit + +# Ambient configured dtype, used only to build inputs. The eager-parity comparisons +# that needed a dtype-dependent tolerance here now live in ``test_gradients.py``, where +# the whole package is pinned to float64 by a fixture rather than reading the global at +# import time -- an import-time read made this file's strictness depend silently on +# ``TORCHREF_DTYPE_FLOAT``. +_F = dtypes.float + + +def _inputs(N=4, R=7, seed=0, dtype=None): + dtype = dtype or _F + torch.manual_seed(seed) + hkl = torch.randint(-3, 4, (R, 3)).to(dtype) + s = torch.rand(R, dtype=dtype) * 0.5 + svec = torch.randn(R, 3, dtype=dtype) * 0.3 + A = torch.rand(N, 5, dtype=dtype) + B = torch.rand(N, 5, dtype=dtype) + 0.5 + return hkl, s, svec, A, B + + +def _leaves(N=4, seed=1, dtype=None): + dtype = dtype or _F + torch.manual_seed(seed) + xyz = torch.rand(N, 3, dtype=dtype, requires_grad=True) + occ = (torch.rand(N, dtype=dtype) * 0.4 + 0.6).requires_grad_() + adp = (torch.rand(N, dtype=dtype) * 10 + 5).requires_grad_() + U = (torch.rand(N, 6, dtype=dtype) * 0.04 + 0.01).requires_grad_() + return xyz, occ, adp, U + + + + +def test_checkpointed_chunking_is_exact(): + # Chunking slices the reflection dimension, so each reflection's atom sum is + # computed within a single chunk -- the result is mathematically identical + # regardless of chunk size. It is not, however, guaranteed bit-exact: the + # phase matmul dispatches to different BLAS kernels for (R,3)x(3,N) vs + # (1,3)x(3,N), and last-ULP rounding there is CPU/BLAS-dependent. Assert + # numerical equivalence, not bitwise equality. + hkl, s, _, A, B = _inputs(R=11, dtype=torch.float64) + xyz, occ, adp, _ = _leaves(dtype=torch.float64) + F_full = D._checkpointed_iso(hkl, s, xyz, occ, adp, A, B, max_memory_gb=None) + F_chunk = D._checkpointed_iso(hkl, s, xyz, occ, adp, A, B, max_memory_gb=1e-7) + assert torch.allclose(F_full, F_chunk, rtol=1e-10, atol=1e-12) + + + + +def test_empty_atoms_returns_zeros(): + hkl, s, _, _, _ = _inputs() + empty = torch.zeros(0, 3, dtype=torch.float64) + z = torch.zeros(0, dtype=torch.float64) + z5 = torch.zeros(0, 5, dtype=torch.float64) + F = ds_iso(hkl, s, empty, z, z, z5, z5) + assert F.shape == (hkl.shape[0],) + assert F.abs().sum().item() == 0.0 + + +def test_integer_hkl_is_accepted_and_lossless(): + """An ``int32`` ``hkl`` must keep working, and must give the float32 answer exactly. + + Miller indices arrive as ``int32`` straight from the MTZ reader + (``torchref/io/mtz.py``), and every DS path casts ``hkl`` itself -- ``_cols_f32`` in + the Triton kernel, ``hkl_c.to(xyz_frac.dtype)`` in the chunk math. So integer input + is not merely tolerated, it is the production case. + + This matters for how the dtype gate may be tightened. The gate's test is + ``t.dtype is torch.float32``, an identity check, so probing ``hkl`` with that rule + would reject the production dtype and disable the Triton path outright. The rule has to + exempt integer tensors, and this test is what fails if it does not. + + Worth recording alongside: casting ``hkl`` is not merely tolerable, it is *free*. + Miller indices are integer-valued in every dtype that holds them (symmetry maps + ``h -> h.R`` with integer ``R``), so the f64->f32 round-trip is bit-exact for any + \\|h\\| < 2**24 and downcasting ``hkl`` alone was measured to change ``F`` by exactly + 0.0. An earlier version of this comment claimed the cast truncated the phase; it does + not, and the gate correspondingly does not probe ``hkl``. + """ + _, s, _, A, B = _inputs(dtype=torch.float64) + xyz, occ, adp, _ = _leaves(dtype=torch.float64) + torch.manual_seed(0) + hkl_int = torch.randint(-3, 4, (7, 3), dtype=torch.int32) + s = s[: hkl_int.shape[0]] + + F_int = ds_iso(hkl_int, s, xyz, occ, adp, A, B) + F_float = ds_iso(hkl_int.to(torch.float64), s, xyz, occ, adp, A, B) + assert torch.equal(F_int, F_float), "casting integer Miller indices must be exact" + + +def test_eager_none_scattering_factors_and_no_batching(): + """``max_memory_gb=None`` with ``scattering_factors=None`` must compute from A/B. + + Two independently optional arguments meeting: with no batching *and* no precomputed + scattering factors, the eager path takes a branch that derives ``f(s)`` from the ITC92 + A/B coefficients inline. Chunking must not change the answer. + + Moved from ``tests/unit/test_kernel_fixes.py``, where it sat among dtype-switch and + NaN-safety tests. It is the only structure-factor test in that file, and this package is + where structure-factor coverage lives. + """ + from torchref.base.direct_summation.dispatch import _eager_aniso, _eager_iso + + g = torch.Generator().manual_seed(0) + N, R = 8, 12 + d = torch.float64 + hkl = torch.randint(-3, 4, (R, 3), generator=g).to(d) + s = torch.rand(R, generator=g, dtype=d) * 0.4 + svec = torch.randn(R, 3, generator=g, dtype=d) * 0.3 + A = torch.rand(N, 5, generator=g, dtype=d) + B = torch.rand(N, 5, generator=g, dtype=d) + 0.5 + xyz = torch.rand(N, 3, generator=g, dtype=d) + occ = torch.rand(N, generator=g, dtype=d) * 0.4 + 0.6 + adp = torch.rand(N, generator=g, dtype=d) * 10 + 5 + U = torch.rand(N, 6, generator=g, dtype=d) * 0.04 + 0.01 + + for tag, fn, geom, third in ( + ("iso", _eager_iso, s, adp), + ("aniso", _eager_aniso, svec, U), + ): + f_none = fn(hkl, geom, xyz, occ, third, A, B, None) + f_batch = fn(hkl, geom, xyz, occ, third, A, B, 2.0) + assert torch.isfinite(f_none).all(), f"{tag}: non-finite F with no batching" + torch.testing.assert_close( + f_none, f_batch, rtol=1e-12, atol=1e-12, + msg=f"{tag}: chunking changed the answer", + ) diff --git a/tests/unit/structure_factor/test_forward.py b/tests/unit/structure_factor/test_forward.py new file mode 100644 index 00000000..0d4caa00 --- /dev/null +++ b/tests/unit/structure_factor/test_forward.py @@ -0,0 +1,521 @@ +"""Forward structure factors: gemmi -> direct summation -> the FFT/splat route. + +Read the package docstring first for the oracle chain and, in particular, for what the +gemmi comparison does and does not prove -- the ITC92 table was generated from gemmi, so +``f(s)`` is shared and cancels. :func:`test_stored_table_matches_gemmi` is the only test +here that constrains the form factors themselves. +""" + +from __future__ import annotations + +import pytest +import torch + +from torchref.model.sf_ds import SfDS + +from . import ( + COS_MIN_DS, + RTOL_DS_F32, + RTOL_DS_F64, + ATOL_TABLE_VS_GEMMI, + COS_MIN, + MAXREL_VS_GEMMI, + RTOL_AMPLITUDE, + RTOL_BACKEND_F32, + RTOL_BACKEND_F64, + RTOL_VS_GEMMI, + SCALE_TOL, + SCALE_TOL_GEMMI, +) +from . import helpers as H +from .conftest import DEVICE_DTYPE_KERNELS, DS_DEVICE_DTYPE_KERNELS + +pytestmark = pytest.mark.unit + + +def _lbl(pin): + """Label for the backend a pin parametrization selected.""" + return "portable" if pin else "default" + +# float32 first: it is the production dtype, so it is the case that matters and the one +# these gates are calibrated on. float64 separates truncation error from float32 noise. +_DTYPES = [torch.float32, torch.float64] + + +def _complex_cos(got, ref): + """Cosine of two complex F vectors, treating them as real 2R-vectors. + + Catches a conjugated result, which an amplitude comparison cannot. + """ + g = torch.view_as_real(got.reshape(-1)).reshape(-1).double() + r = torch.view_as_real(ref.reshape(-1)).reshape(-1).double() + return float((g @ r) / (g.norm() * r.norm()).clamp_min(1e-30)) + + +def _report(tag, got, ref): + """Print the three metrics, so a widened tolerance always has a number behind it.""" + r, m, k = H.rel_l2(got, ref), H.max_rel(got, ref), H.best_fit_scale(got, ref) + print(f" {tag:34s} relL2 {r:.3e} maxrel {m:.3e} scale {k:.9f}") + return r, m, k + + +# --------------------------------------------------------------------------- +# The one place gemmi is an independent reference for f(s) +# --------------------------------------------------------------------------- +def test_stored_table_matches_gemmi(): + """The stored ITC92 ``.pt`` must still equal gemmi's live table. + + ``torchref/data/itc92_scattering_factors.pt`` is generated from gemmi by + ``torchref/scripts/generate_scattering_table.py``, then used at runtime so gemmi is + not a runtime dependency. Nothing re-checked the two after generation, so a + corrupted regeneration or a gemmi release that revised coefficients would go + unnoticed -- and would be invisible to every other test in this package, because + those compare torchref against gemmi using the *same* coefficients on both sides. + + Also pins the storage convention: torchref keeps gemmi's four Gaussians plus the + constant ``c`` folded in as a fifth with ``B = 0``, so ``exp(0) = 1`` reproduces the + additive term. + + Compared in **float32**, and required to be bit-exact there. The stored table is + float32 (``(104, 5)``), while gemmi's Python ``it92`` returns doubles -- so a + float64 comparison can only ever agree to float32 rounding, measured at 5.9e-08 + worst case over Z=1..103, which is half of float32 eps. Asserting exact equality + after rounding gemmi's doubles to float32 is the stronger statement: it permits the + intended precision loss and nothing else, where a 1e-7 rtol would also permit a + genuinely wrong coefficient that happened to land nearby. + + Note this caps the oracle: asking ``get_scattering_params_by_z`` for float64 widens + float32 values rather than recovering the doubles, so ``f(s)`` in the "float64" + oracle carries ~6e-8 relative error. That is the floor behind ``RTOL_VS_GEMMI`` and + it is the correct trade for a float32 production path. + """ + gemmi = pytest.importorskip("gemmi") + from torchref.base.scattering.scattering_table import get_scattering_params_by_z + + checked = 0 + for z in range(1, 104): + element = gemmi.Element(z) + coef = element.it92 + if coef is None: + continue + A, B = get_scattering_params_by_z(torch.tensor([z]), dtype=torch.float32) + assert A.shape == (1, 5) and B.shape == (1, 5), f"Z={z} unexpected table shape" + + want_a = torch.tensor(list(coef.a), dtype=torch.float32) + want_b = torch.tensor(list(coef.b), dtype=torch.float32) + want_c = torch.tensor(coef.c, dtype=torch.float32) + + assert torch.equal(A[0, :4], want_a), ( + f"Z={z} ({element.name}): 'a' drifted from gemmi -- " + f"stored {A[0, :4].tolist()} vs gemmi {want_a.tolist()}" + ) + assert torch.equal(B[0, :4], want_b), ( + f"Z={z} ({element.name}): 'b' drifted from gemmi -- " + f"stored {B[0, :4].tolist()} vs gemmi {want_b.tolist()}" + ) + assert torch.equal(A[0, 4], want_c), ( + f"Z={z} ({element.name}): constant term c not in slot 4 -- " + f"stored {float(A[0, 4])} vs gemmi {float(want_c)}" + ) + assert float(B[0, 4]) == 0.0, ( + f"Z={z} ({element.name}): slot 4 must have B=0 so exp(0)=1 gives +c" + ) + checked += 1 + + assert checked >= 90, f"only {checked} elements checked -- table lookup is not working" + + +def test_ls_target_is_phase_blind(scene_small): + """An amplitude target and its gradients are invariant under conjugating ``F``. + + This is why phase convention is gated on forward complex values -- here and against + gemmi -- and *not* through the derivative tests. Since + + d|F|/dx = Re(F* dF/dx) / |F| + + sending ``F -> F*`` and ``dF/dx -> (dF/dx)*`` leaves ``Re(F* dF/dx)`` unchanged, so a + global phase convention is an unobservable gauge choice for such a target. An earlier + draft tried to catch a conjugated kernel with a functional linear in ``F.imag``; that + cannot work for a physical target and cost a 40x conditioning artefact in the relative + error (2.7e-01 vs ~6e-03 for the same comparison). See the note in ``helpers.py``. + + Phase still matters -- anomalous scattering and difference maps are phase-sensitive, + and the two routes must agree with each other -- which is exactly what + :func:`test_ds_matches_gemmi_iso_p1` and the complex-cosine assertion in + :func:`test_fft_matches_ds_amplitudes` cover. + """ + s = scene_small + with torch.no_grad(): + F = H.ds_iso_oracle(s) + obs = H.synthetic_obs(F) + + def target(x, conjugate): + Fx = H.ds_iso_oracle(s, x, s.occ, s.adp) + return H.ls_target(Fx.conj() if conjugate else Fx, obs) + + x1 = s.xyz.clone().requires_grad_(True) + x2 = s.xyz.clone().requires_grad_(True) + v1, v2 = target(x1, False), target(x2, True) + (g1,) = torch.autograd.grad(v1, x1) + (g2,) = torch.autograd.grad(v2, x2) + + assert v1.item() == pytest.approx(v2.item(), rel=1e-12), ( + "conjugating F changed the amplitude target value" + ) + assert H.rel_l2(g2, g1) < 1e-12, ( + "conjugating F changed the amplitude target's gradient, so this target is not " + "phase-blind after all -- the derivative tests could then gate phase convention " + "and the reasoning in helpers.py needs revisiting" + ) + + +# --------------------------------------------------------------------------- +# gemmi -> DS +# --------------------------------------------------------------------------- +def test_ds_matches_gemmi_iso_p1(gemmi_iso_p1): + """Isotropic P1. Constrains the exponential sum, the 2*pi factors and the B + convention, with symmetry taken out of the picture.""" + scene, structure = gemmi_iso_p1 + F_gemmi = H.gemmi_sf(structure, scene.hkl_list) + with torch.no_grad(): + F_ds = H.ds_iso_oracle(scene) + + print(f"\n{scene.n_atoms} atoms, {scene.n_refl} reflections, P1") + rel, mrel, scale = _report("DS vs gemmi (iso, P1)", F_ds, F_gemmi) + assert rel < RTOL_VS_GEMMI + assert mrel < MAXREL_VS_GEMMI + assert abs(scale - 1.0) < SCALE_TOL_GEMMI + + +def test_ds_matches_gemmi_aniso_p1(gemmi_aniso_p1): + """Anisotropic P1 on a structure where every atom carries ANISOU. + + Compares complex ``F``, not ``|F|``: an amplitude-only check passes a phase-sign + error, and the aniso Debye-Waller factor is real, so a sign slip elsewhere would + hide behind it. + """ + scene, structure = gemmi_aniso_p1 + assert (scene.u6.abs().sum(dim=1) > 0).all(), ( + "every atom must carry ANISOU or this test degenerates to the isotropic case" + ) + F_gemmi = H.gemmi_sf(structure, scene.hkl_list) + with torch.no_grad(): + F_ds = H.ds_aniso_oracle(scene) + + print(f"\n{scene.n_atoms} ANISOU atoms, {scene.n_refl} reflections, P1") + rel, mrel, scale = _report("DS vs gemmi (aniso, P1)", F_ds, F_gemmi) + assert rel < RTOL_VS_GEMMI + assert mrel < MAXREL_VS_GEMMI + assert abs(scale - 1.0) < SCALE_TOL_GEMMI + + +@pytest.mark.parametrize( + "perm,label", + [((0, 1, 2, 4, 3, 5), "U12<->U13"), ((0, 1, 2, 5, 4, 3), "reversed off-diagonals")], +) +def test_aniso_ordering_is_actually_constrained(gemmi_aniso_p1, perm, label): + """Non-vacuity guard for the test above. + + A tight tolerance is only meaningful if a wrong convention would breach it. Permuting + the ADP off-diagonals must move the residual far outside ``RTOL_VS_GEMMI`` -- measured + 1.35e-02 for the ``U12<->U13`` swap and 6.08e-03 for the reversal, five orders of + magnitude clear. Without this, ``[U11,U22,U33,U12,U13,U23]`` would be an assumption + rather than a verified fact. + """ + scene, structure = gemmi_aniso_p1 + F_gemmi = H.gemmi_sf(structure, scene.hkl_list) + with torch.no_grad(): + F_perm = H.ds_aniso_oracle(scene, u6=scene.u6[:, list(perm)]) + + rel = H.rel_l2(F_perm, F_gemmi) + print(f"\n permuted {label:24s} relL2 {rel:.3e} (gate is {RTOL_VS_GEMMI:.0e})") + assert rel > 100 * RTOL_VS_GEMMI, ( + f"permuting {label} changed the result by only {rel:.2e}; the aniso comparison " + "does not constrain the off-diagonal ordering" + ) + + +def test_sfds_matches_gemmi_with_symmetry(gemmi_iso_symmetry): + """Symmetry algebra against a second implementation, in a hexagonal group. + + ``P 65 2 2`` is chosen because ``h.R`` and ``R.h`` agree for symmetric rotation + matrices and differ for trigonal/hexagonal ones -- so an orthorhombic or tetragonal + entry could not detect the transpose error that + ``tests/unit/symmetry/test_hkl_symmetry_gemmi.py`` documents. + + This is also the only symmetric comparison in the package. A DS-vs-FFT check cannot + validate symmetry, because both routes call the same + ``compute_symmetry_equivalent_hkls`` / ``compute_translation_phases`` and the shared + algebra cancels; gemmi does not share it. + """ + scene, structure = gemmi_iso_symmetry + assert len(structure.cell.images) > 0, "structure was not set up with symmetry" + + F_gemmi = H.gemmi_sf(structure, scene.hkl_list) + ds = SfDS( + cell=scene.cell, + spacegroup=scene.spacegroup, + dtype_float=torch.float64, + device=torch.device("cpu"), + ) + with torch.no_grad(): + F_sym, _ = ds.compute_structure_factors( + scene.hkl, scene.xyz, scene.adp, scene.occ, scene.A, scene.B, + apply_symmetry=True, + ) + + print(f"\n{scene.spacegroup}, {len(structure.cell.images) + 1} operations") + rel, mrel, scale = _report("SfDS(sym) vs gemmi", F_sym, F_gemmi) + assert rel < RTOL_VS_GEMMI + assert mrel < MAXREL_VS_GEMMI + assert abs(scale - 1.0) < SCALE_TOL_GEMMI + + +# --------------------------------------------------------------------------- +# DS -> FFT +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("pin", [False, True], ids=["default", "portable"]) +@pytest.mark.parametrize("dtype", _DTYPES, ids=["f32", "f64"]) +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +def test_fft_matches_ds_amplitudes(scene_fine, oracle_fine, kind, dtype, pin): + """The headline gate: the map route reproduces the analytic answer. + + The DS reference stays in float64 whatever the candidate's dtype -- an oracle should + not inherit the precision of the thing it is judging. + """ + aniso = kind == "aniso" + sf = H.sf_fft_for(scene_fine, dtype) + with H.maybe_portable(pin), torch.no_grad(): + F_fft = H.fft_sf(scene_fine, sf, aniso=aniso).to(torch.complex128) + F_ds = oracle_fine[f"{kind}_F"] + + print(f"\ngrid {tuple(int(v) for v in sf.gridsize)}, {kind}, {dtype}, {_lbl(pin)}") + rel, mrel, scale = _report(f"FFT vs DS ({kind})", F_fft, F_ds) + assert rel < RTOL_AMPLITUDE, f"{kind}/{dtype}/{_lbl(pin)}: rel L2 {rel:.3e}" + cos = float( + (F_fft.conj() * F_ds).real.sum() + / (F_fft.abs().norm() * F_ds.abs().norm()).clamp_min(1e-30) + ) + assert cos > COS_MIN, f"phase disagreement: cos {cos:.6f}" + + +def test_fft_absolute_scale(scene_fine, oracle_fine): + """Absolute scale, which nothing in the suite checked before. + + The FFT route picks up a voxel-volume factor from ``ifft`` + (``torchref/base/fourier/fft.py``) while DS is a bare atom sum. A factor of ``V``, + of ``N`` voxels, or of ``V/N`` would leave every ratio-, cosine- and parity-based + gate in the repo perfectly happy -- correlation stays 1.0 under a global rescale. + """ + sf = H.sf_fft_for(scene_fine) + with torch.no_grad(): + F_fft = H.fft_sf(scene_fine, sf) + scale = H.best_fit_scale(F_fft, oracle_fine["iso_F"]) + print(f"\n best-fit scale {scale:.9f}") + assert abs(scale - 1.0) < SCALE_TOL, ( + f"absolute scale off by {abs(scale - 1.0):.3e}; suspect a voxel-volume or " + f"grid-size normalization factor" + ) + + +def test_fft_matches_ds_monoclinic(scene_monoclinic, oracle_monoclinic): + """beta = 115 deg. A truncation that is spherical in fractional rather than + Cartesian space, or centred on a grid node rather than the atom, diverges most in a + strongly non-orthogonal metric -- measured 5.0e-3 rel L2 for the node-centred + variant at this angle, against 9.9e-4 orthorhombic.""" + sf = H.sf_fft_for(scene_monoclinic) + with torch.no_grad(): + F_fft = H.fft_sf(scene_monoclinic, sf) + rel, _, _ = _report("FFT vs DS (beta=115)", F_fft, oracle_monoclinic["iso_F"]) + assert rel < RTOL_AMPLITUDE + + +@pytest.mark.parametrize("dtype", _DTYPES, ids=["f32", "f64"]) +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +def test_backend_parity_auto_vs_eager(scene_fine, kind, dtype): + """AUTO and EAGER implement the same truncation contract, so they must agree to + float noise -- far tighter than either agrees with DS. This separates a kernel bug + from the shared truncation error that ``test_fft_matches_ds_amplitudes`` bounds.""" + aniso = kind == "aniso" + sf = H.sf_fft_for(scene_fine, dtype) + with torch.no_grad(): + with H.maybe_portable(False): + auto = H.fft_sf(scene_fine, sf, aniso=aniso).to(torch.complex128) + with H.maybe_portable(True): + eager = H.fft_sf(scene_fine, sf, aniso=aniso).to(torch.complex128) + tol = RTOL_BACKEND_F32 if dtype is torch.float32 else RTOL_BACKEND_F64 + rel = H.rel_l2(auto, eager) + print(f"\n AUTO vs EAGER ({kind}, {dtype}): relL2 {rel:.3e} gate {tol:.0e}") + assert rel < tol + + +# --------------------------------------------------------------------------- +# Truncation budget +# --------------------------------------------------------------------------- +def test_nsigma_reduces_truncation_error(scene_fine, oracle_fine, sigma_cutoff): + """Raising the cutoff must monotonically improve agreement with DS. + + Asserted as a trend rather than an absolute number, for two reasons: it puts a test + behind the N_sigma claim at ``torchref/config.py:196`` which until now existed only + as a comment, and it stays valid whichever default the cutoff ends up at, so this + test does not have to be retuned alongside that decision. + """ + # Explicitly finer than production: at d_min/3 the cutoff is not the binding + # constraint (5.89e-3 -> 5.86e-3 then flat), so a monotonicity assertion there + # would be measuring grid sampling. See GRID_FINENESS in helpers.py. + sf = H.sf_fft_for(scene_fine, fineness=1.6) + residuals = [] + for n_sigma in (2.5, 3.0, 3.5, 4.5): + sigma_cutoff(n_sigma) + with torch.no_grad(): + F_fft = H.fft_sf(scene_fine, sf) + residuals.append((n_sigma, H.rel_l2(F_fft, oracle_fine["iso_F"]))) + + print() + for n_sigma, rel in residuals: + print(f" n_sigma {n_sigma:4.1f} relL2 vs DS {rel:.3e}") + + values = [r for _, r in residuals] + assert values == sorted(values, reverse=True), ( + f"truncation error is not monotone in n_sigma: {residuals}" + ) + assert values[0] > values[-1] * 2, ( + "tightening the cutoff barely changed the residual, so this scene is " + "sampling-limited rather than truncation-limited and does not test the cutoff" + ) + + +def test_coarse_grid_is_measurably_worse(scene_coarse, scene_fine, oracle_fine): + """Non-vacuity guard for ``RTOL_AMPLITUDE``. + + If an under-sampled grid also passed the gate, the gate would be measuring nothing. + This pins the sampling-dominated regime separately so a regression that degraded + grid sampling cannot hide behind the fine-grid tolerance. + """ + # Bare Nyquist (oversampling 2 rather than production's 3). + sf_coarse = H.sf_fft_for(scene_fine, fineness=2.0 / 3.0) + with torch.no_grad(): + F_coarse = H.fft_sf(scene_fine, sf_coarse) + rel = H.rel_l2(F_coarse, oracle_fine["iso_F"]) + print( + f"\n coarse grid {tuple(int(v) for v in sf_coarse.gridsize)}: " + f"relL2 {rel:.3e} (fine-grid gate is {RTOL_AMPLITUDE:.0e})" + ) + assert rel > RTOL_AMPLITUDE, ( + "an under-sampled grid passes the amplitude gate, so the gate is not " + "constraining grid sampling at all" + ) + + +# --------------------------------------------------------------------------- +# Every production kernel, on every device it ships on +# --------------------------------------------------------------------------- +# These call each splat kernel **directly** rather than through +# ``build_electron_density``. Under the default a failed accelerator +# kernel silently falls back to the portable splat, so a dispatch-driven test can pass +# while measuring a different kernel than the one it names; a direct call settles that by +# construction. The dispatch ladder is tested separately in ``test_dispatch.py``. +# +# The oracle stays CPU float64 whatever the candidate is -- a reference must not inherit +# the precision of the thing it judges. +# +# Measured when written (60-atom synthetic scene, worst case over the whole matrix): +# +# amplitude rel L2 5.68e-03 gradient rel 9.51e-02 +# amplitude cos 0.999990 gradient cos 0.995544 +# +# and, notably, **every device agrees to the printed digits**: Metal float32, CPU fused +# float32/float64 and the portable splat on either device all land on the same residual +# against the oracle. That is the truncation-contract standardization showing up as a +# measurement -- what is left is shared discretization error, not per-kernel divergence. +# Hence no accelerator-specific tolerances: the CPU constants already bound it. + + +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +@pytest.mark.parametrize("device,dtype,kernel", DEVICE_DTYPE_KERNELS) +def test_kernel_matches_ds_amplitudes(scene_fine, oracle_fine, device, dtype, kernel, kind): + """Forward amplitudes from one production kernel against the DS oracle.""" + aniso = kind == "aniso" + scene = scene_fine.to(device=device, dtype=dtype) + with torch.no_grad(): + got = H.density_to_F( + scene, H.splat_direct(scene, kernel, aniso=aniso) + ).cpu().to(torch.complex128) + ref = oracle_fine[f"{kind}_F"] + + rel, cos = H.rel_l2(got, ref), _complex_cos(got, ref) + print(f"\n {device.type}/{dtype}/{kernel}/{kind}: relL2 {rel:.3e} cos {cos:.8f}") + assert rel < RTOL_AMPLITUDE, f"{device.type}/{dtype}/{kernel}/{kind}: rel {rel:.3e}" + assert cos > COS_MIN, f"{device.type}/{dtype}/{kernel}/{kind}: cos {cos:.6f}" + + +@pytest.mark.parametrize("device,dtype,kernel", DEVICE_DTYPE_KERNELS) +def test_kernel_absolute_scale(scene_fine, oracle_fine, device, dtype, kernel): + """Absolute scale per kernel. + + Runs on every backend because the voxel-volume factor comes from ``ifft``, which is + shared, but the *density* each kernel accumulates is not -- a kernel that normalized + its Gaussian differently would show up only here. Correlation and cosine are both + invariant to a global rescale. + """ + scene = scene_fine.to(device=device, dtype=dtype) + with torch.no_grad(): + got = H.density_to_F(scene, H.splat_direct(scene, kernel)).cpu().to(torch.complex128) + scale = H.best_fit_scale(got, oracle_fine["iso_F"]) + print(f"\n {device.type}/{dtype}/{kernel}: scale {scale:.9f}") + assert abs(scale - 1.0) < SCALE_TOL, ( + f"{device.type}/{dtype}/{kernel}: absolute scale off by {abs(scale - 1.0):.3e}" + ) + + +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +@pytest.mark.parametrize("device,dtype,kernel", DEVICE_DTYPE_KERNELS) +def test_kernels_agree_with_each_other(scene_fine, device, dtype, kernel, kind): + """Every kernel on a given ``(device, dtype)`` must agree with the portable splat. + + Complements the oracle comparison rather than duplicating it. The oracle gate is + ``1e-2``, loose enough to hide a kernel that is subtly wrong in the same direction as + the discretization error; this one is tight, because two kernels implementing the same + truncation contract on the same inputs have nothing to disagree about beyond float + accumulation order. + """ + if kernel == "portable": + pytest.skip("portable is the reference for this comparison") + aniso = kind == "aniso" + scene = scene_fine.to(device=device, dtype=dtype) + with torch.no_grad(): + got = H.splat_direct(scene, kernel, aniso=aniso) + ref = H.splat_direct(scene, "portable", aniso=aniso) + tol = RTOL_BACKEND_F32 if dtype is torch.float32 else RTOL_BACKEND_F64 + rel = H.rel_l2(got.cpu().to(torch.float64), ref.cpu().to(torch.float64)) + print(f"\n {device.type}/{dtype}/{kernel} vs portable ({kind}): relL2 {rel:.3e}") + assert rel < tol, f"{device.type}/{dtype}/{kernel}/{kind}: vs portable rel {rel:.3e}" + + +# --------------------------------------------------------------------------- +# Direct-summation kernels vs the eager oracle +# --------------------------------------------------------------------------- +# Restores coverage lost when ``tests/integration/test_ds_triton_vs_eager.py`` was deleted, +# and extends it: that file only ran on CUDA, so ``_checkpointed_*`` -- what every CPU, MPS +# and float64 call actually executes -- was compared against eager on CPU only, and DS on +# MPS had no coverage at all. +# +# No grid and no truncation on either side, so the only source of disagreement is float +# arithmetic and these gate near precision. See ``RTOL_DS_F32`` / ``RTOL_DS_F64``. + + +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +@pytest.mark.parametrize("device,dtype,kernel", DS_DEVICE_DTYPE_KERNELS) +def test_ds_kernel_matches_eager_oracle(scene_fine, device, dtype, kernel, kind): + """Forward ``F(hkl)`` from one DS kernel against the pure-torch eager oracle.""" + aniso = kind == "aniso" + scene = scene_fine.to(device=device, dtype=dtype) + with torch.no_grad(): + got = H.ds_direct(scene, kernel, aniso=aniso).cpu().to(torch.complex128) + ref = H.ds_direct(scene_fine, "eager", aniso=aniso) + + tol = RTOL_DS_F32 if dtype is torch.float32 else RTOL_DS_F64 + rel, cos = H.rel_l2(got, ref), _complex_cos(got, ref) + print(f"\n {device.type}/{dtype}/{kernel}/{kind}: relL2 {rel:.3e} cos {cos:.9f}") + assert rel < tol, f"{device.type}/{dtype}/{kernel}/{kind}: rel {rel:.3e}" + assert cos > COS_MIN_DS, f"{device.type}/{dtype}/{kernel}/{kind}: cos {cos:.7f}" diff --git a/tests/unit/structure_factor/test_gradients.py b/tests/unit/structure_factor/test_gradients.py new file mode 100644 index 00000000..70604fd4 --- /dev/null +++ b/tests/unit/structure_factor/test_gradients.py @@ -0,0 +1,416 @@ +"""First-order gradients: FD -> direct summation -> the FFT/splat route. + +Three links, in order: + +1. ``gradcheck`` proves the eager DS path's gradients against finite differences. DS is + analytic, so FD converges against it and this gate is tight. +2. The *production* DS path (``_checkpointed_*``, what ``ds_iso`` and ``SfDS`` actually + call) is tied to the eager oracle. This link is load-bearing: without it, "the oracle + is correct" would say nothing about the code that ships. +3. The FFT/splat route is compared against the oracle, at the loose gate the truncation + and sampling error actually warrants. +""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from tests.helpers.grad_asserts import ( + assert_grads_agree, + cosine_similarity, + gradnorm_ratio, + rel_error, +) +from torchref.base.direct_summation.dispatch import ( + _checkpointed_aniso, + _checkpointed_iso, + _eager_aniso, + _eager_iso, +) + +from . import ( + COS_MIN_DS, + RTOL_DS_F32, + RTOL_DS_F64, + ATOL_GRADCHECK, + COS_MIN_GRADIENT, + COS_MIN_GRADIENT_SYNTHETIC, + EPS_GRADCHECK, + RTOL_BACKEND_GRAD_F32, + RTOL_BACKEND_GRAD_F64, + RTOL_GRADCHECK, + RTOL_GRADIENT, + RTOL_GRADIENT_SYNTHETIC, +) +from . import helpers as H +from .conftest import DEVICE_DTYPE_KERNELS, DS_DEVICE_DTYPE_KERNELS + +pytestmark = pytest.mark.unit + + +def _lbl(pin): + """Label for the backend a pin parametrization selected.""" + return "portable" if pin else "default" + +_DTYPES = [torch.float32, torch.float64] # float32 first: production dtype +_LEAF_NAMES = ("xyz", "occ", "adp_or_u") + + +# --------------------------------------------------------------------------- +# 1. FD -> DS +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +def test_ds_gradcheck(scene_small, kind): + """The oracle's own gradients, against finite differences. + + Run on the deliberately tiny ``scene_small`` because ``gradcheck`` costs + O(n_params) forward evaluations -- 430 ms at 10 atoms, 75 s at 60. The property being + proved (that eager DS differentiates correctly) is scene-independent, so there is + nothing to gain from a larger scene and a great deal of wall-clock to lose. + """ + s = scene_small + if kind == "iso": + fn = lambda x, o, a: _eager_iso(s.hkl, s.s, x, o, a, s.A, s.B, None) # noqa: E731 + third = s.adp + else: + fn = lambda x, o, u: _eager_aniso( # noqa: E731 + s.hkl, s.s_vec, x, o, u, s.A, s.B, None + ) + third = s.u6 + + args = tuple( + t.clone().requires_grad_(True) for t in (s.xyz_frac, s.occ, third) + ) + assert torch.autograd.gradcheck( + fn, args, eps=EPS_GRADCHECK, atol=ATOL_GRADCHECK, rtol=RTOL_GRADCHECK + ) + + +# --------------------------------------------------------------------------- +# 2. The production DS path matches the oracle +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("chunked", [False, True], ids=["one_chunk", "chunked"]) +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +def test_checkpointed_matches_eager_oracle(scene_fine, kind, chunked): + """``_checkpointed_*`` -- reached by every ``ds_iso``/``ds_aniso``/``SfDS`` call on + CPU, MPS or float64 -- must reproduce the eager oracle in both value and gradient. + + This is the link that makes ``_eager_*`` a legitimate stand-in for the shipping code. + ``_CheckpointedSF`` recomputes each chunk under ``enable_grad`` in its backward, so + the gradient path is genuinely different from eager even though the maths is the + same; ``chunked`` forces multiple chunks via a tiny memory budget so that recompute + loop is actually exercised rather than degenerating to one pass. + + Measured agreement is 2.3e-16, so this is gated near machine precision. It is + deliberately far tighter than anything involving the map route: any drift here is a + bug, not an approximation. + """ + s = scene_fine + max_mem = 1e-7 if chunked else None + if kind == "iso": + eager = lambda x, o, a: _eager_iso(s.hkl, s.s, x, o, a, s.A, s.B, max_mem) # noqa: E731 + prod = lambda x, o, a: _checkpointed_iso( # noqa: E731 + s.hkl, s.s, x, o, a, s.A, s.B, max_mem + ) + third = s.adp + else: + eager = lambda x, o, u: _eager_aniso( # noqa: E731 + s.hkl, s.s_vec, x, o, u, s.A, s.B, max_mem + ) + prod = lambda x, o, u: _checkpointed_aniso( # noqa: E731 + s.hkl, s.s_vec, x, o, u, s.A, s.B, max_mem + ) + third = s.u6 + + with torch.no_grad(): + obs = H.synthetic_obs(eager(s.xyz_frac, s.occ, third)) + + def run(fn): + leaves = tuple( + t.clone().requires_grad_(True) for t in (s.xyz_frac, s.occ, third) + ) + F = fn(*leaves) + grads = torch.autograd.grad(H.ls_target(F, obs), leaves) + return F.detach(), grads + + F_eager, g_eager = run(eager) + F_prod, g_prod = run(prod) + + assert rel_error(F_prod, F_eager) < 1e-13, "checkpointed forward drifted from eager" + assert_grads_agree( + dict(zip(_LEAF_NAMES, g_prod)), + dict(zip(_LEAF_NAMES, g_eager)), + min_cos=1.0 - 1e-12, + ratio_tol=1e-10, + ctx=f"{kind}/{'chunked' if chunked else 'one_chunk'} ", + ) + + +def test_public_ds_api_matches_oracle(scene_fine): + """The same check one level up, through ``ds_iso``'s own dispatch. + + Guards against the dispatch layer -- not the kernel -- introducing a discrepancy: + a wrong ``max_memory_gb`` default, or a dispatch table that silently routes + somewhere unintended on CPU. + """ + from torchref.base.direct_summation.dispatch import ds_iso + + s = scene_fine + with torch.no_grad(): + obs = H.synthetic_obs(_eager_iso(s.hkl, s.s, s.xyz_frac, s.occ, s.adp, s.A, s.B, None)) + + leaves = tuple(t.clone().requires_grad_(True) for t in (s.xyz_frac, s.occ, s.adp)) + F = ds_iso(s.hkl, s.s, *leaves, s.A, s.B) + g_public = torch.autograd.grad(H.ls_target(F, obs), leaves) + + leaves2 = tuple(t.clone().requires_grad_(True) for t in (s.xyz_frac, s.occ, s.adp)) + F2 = _eager_iso(s.hkl, s.s, *leaves2, s.A, s.B, None) + g_oracle = torch.autograd.grad(H.ls_target(F2, obs), leaves2) + + assert rel_error(F, F2) < 1e-13 + assert_grads_agree( + dict(zip(_LEAF_NAMES, g_public)), + dict(zip(_LEAF_NAMES, g_oracle)), + min_cos=1.0 - 1e-12, + ratio_tol=1e-10, + ctx="ds_iso ", + ) + + +# --------------------------------------------------------------------------- +# 3. DS -> FFT +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("pin", [False, True], ids=["default", "portable"]) +@pytest.mark.parametrize("dtype", _DTYPES, ids=["f32", "f64"]) +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +def test_fft_gradients_match_ds(scene_fine, oracle_fine, kind, dtype, pin): + """Gradients of the map route against the analytic oracle, for all three leaves. + + ``occ`` is included deliberately. Before this package, occupancy gradients through + the FFT route were untested at any order -- the previous FFT gradient test + parametrized only ``xyz`` and ``adp`` -- and ``occ`` is the one leaf whose gradient a + scene with ``occ == 1`` everywhere cannot constrain, which is why + :func:`helpers.synthetic_scene` never uses unit occupancy. + + Gradients are roughly 10x more truncation-sensitive than amplitudes, because each + ``xyz`` derivative brings a factor ~``2*pi*h``. Hence ``RTOL_GRADIENT`` an order of + magnitude above ``RTOL_AMPLITUDE`` rather than the same number. + """ + aniso = kind == "aniso" + sf = H.sf_fft_for(scene_fine, dtype) + leaves = scene_fine.leaves(aniso=aniso) + with H.maybe_portable(pin): + F = H.fft_sf(scene_fine, sf, *leaves, aniso=aniso) + got = torch.autograd.grad(H.ls_target(F, oracle_fine[f"{kind}_obs"]), leaves) + ref = oracle_fine[f"{kind}_grads"] + + gate, cos_gate = RTOL_GRADIENT_SYNTHETIC, COS_MIN_GRADIENT_SYNTHETIC + print(f"\n{kind}, {dtype}, {_lbl(pin)} (rel gate {gate:.0e}, cos gate {cos_gate})") + for name, g, r in zip(_LEAF_NAMES, got, ref): + rel, cos = rel_error(g, r), cosine_similarity(g, r) + print(f" {name:10s} rel {rel:.3e} cos {cos:.8f} ratio {gradnorm_ratio(g, r):.4f}") + assert rel < gate, f"{kind}/{dtype}/{_lbl(pin)} {name}: rel {rel:.3e}" + assert cos > cos_gate, f"{kind}/{dtype}/{_lbl(pin)} {name}: cos {cos:.6f}" + + +@pytest.mark.parametrize("dtype", _DTYPES, ids=["f32", "f64"]) +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +def test_backend_parity_of_gradients(scene_fine, oracle_fine, kind, dtype): + """AUTO and EAGER share one truncation contract, so their gradients must agree far + more tightly than either agrees with DS. Separates a kernel-specific backward bug + from the shared truncation residual.""" + aniso = kind == "aniso" + sf = H.sf_fft_for(scene_fine, dtype) + + def run(pin): + leaves = scene_fine.leaves(aniso=aniso) + with H.maybe_portable(pin): + F = H.fft_sf(scene_fine, sf, *leaves, aniso=aniso) + return torch.autograd.grad(H.ls_target(F, oracle_fine[f"{kind}_obs"]), leaves) + + auto, eager = run(False), run(True) + tol = RTOL_BACKEND_GRAD_F32 if dtype is torch.float32 else RTOL_BACKEND_GRAD_F64 + for name, a, e in zip(_LEAF_NAMES, auto, eager): + rel = rel_error(a, e) + print(f" {kind}/{dtype} {name:10s} AUTO vs EAGER rel {rel:.3e} (gate {tol:.0e})") + assert rel < tol, f"{kind}/{dtype} {name}: AUTO vs EAGER rel {rel:.3e}" + + +def test_occupancy_gradient_is_not_degenerate(scene_fine, oracle_fine): + """Non-vacuity guard: the ``occ`` gradient must be non-trivial and atom-dependent. + + An all-equal or all-zero ``occ`` gradient would satisfy the cosine and relative gates + above while carrying no information, which is exactly what a scene with uniform + occupancy produces. + """ + g_occ = oracle_fine["iso_grads"][1] + assert g_occ.abs().min() > 0, "some occupancy gradient is exactly zero" + spread = float(g_occ.std() / g_occ.abs().mean()) + print(f"\n occ gradient relative spread across atoms: {spread:.3f}") + assert spread > 0.05, ( + "occupancy gradients are nearly identical across atoms, so the comparison " + "cannot distinguish a correct per-atom gradient from a constant" + ) + + +# --------------------------------------------------------------------------- +# 4. The authoritative gate: a real structure +# --------------------------------------------------------------------------- +def test_fft_gradients_match_ds_real_structure(gemmi_aniso_grad, oracle_aniso_grad): + """Gradients on 7L84 -- 1209 atoms, all with ANISOU -- at production sampling. + + This is the gate that states something about production; the synthetic sweeps above + exist for dtype/backend/cell coverage and are calibrated loosely because small scenes + overstate the discretization residual (see ``__init__.py``). + + Run in float32, the production dtype, against the float64 oracle. Nothing here is + tuned: the numbers came out at 1.35e-02 (xyz) and 2.62e-02 (U) against the 1e-01 gate + you specified, so there is ~4x headroom without the gate being vacuous -- the + under-sampled control in ``test_forward.py`` shows what breaching it looks like. + """ + scene, _ = gemmi_aniso_grad + sf = H.sf_fft_for(scene, torch.float32) + leaves = scene.leaves(aniso=True) + got = torch.autograd.grad( + H.ls_target(H.fft_sf(scene, sf, *leaves, aniso=True), oracle_aniso_grad["aniso_obs"]), + leaves, + ) + ref = oracle_aniso_grad["aniso_grads"] + + print(f"\n7L84 P1, {scene.n_atoms} ANISOU atoms, {scene.n_refl} refl, float32") + print(f" grid {tuple(int(v) for v in sf.gridsize)}") + for name, g, r in zip(_LEAF_NAMES, got, ref): + rel, cos, ratio = rel_error(g, r), cosine_similarity(g, r), gradnorm_ratio(g, r) + print(f" {name:10s} rel {rel:.3e} cos {cos:.8f} ratio {ratio:.4f}") + assert rel < RTOL_GRADIENT, f"7L84 {name}: rel {rel:.3e}" + assert cos > COS_MIN_GRADIENT, f"7L84 {name}: cos {cos:.6f}" + + # The deviatoric part specifically -- an earlier revision of this package recorded a + # 54% bias here from a 10-atom synthetic scene. On a real structure it is not present, + # and this assertion is what keeps that claim from creeping back. + dev = lambda g: g[:, :3] - g[:, :3].mean(dim=1, keepdim=True) # noqa: E731 + dev_ratio = gradnorm_ratio(dev(got[2]), dev(ref[2])) + print(f" deviatoric U ratio {dev_ratio:.4f}") + assert abs(dev_ratio - 1.0) < 0.05, ( + f"deviatoric anisotropic ADP gradient is biased by {abs(dev_ratio - 1.0):.1%} on a " + "real structure. That would justify revisiting the grid oversampling or a B_extra " + "smearing correction; it was measured at 0.9995 when this test was written." + ) + + + +# --------------------------------------------------------------------------- +# 5. Every production kernel, on every device it ships on +# --------------------------------------------------------------------------- +# Direct kernel calls, not dispatch -- see the note in ``test_forward.py``. Worst case +# over the whole matrix when written: gradient rel 9.51e-02, cos 0.995544, and every +# device agreeing to the printed digits, so the CPU constants bound the accelerators too. + + +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +@pytest.mark.parametrize("device,dtype,kernel", DEVICE_DTYPE_KERNELS) +def test_kernel_gradients_match_ds(scene_fine, oracle_fine, device, dtype, kernel, kind): + """Gradients from one production kernel against the DS oracle, all three leaves. + + The candidate's leaves live on the target device in the target dtype, while the + oracle's are CPU float64. Both describe the same structure because + :meth:`Scene.to` moves and casts a single source scene, and the comparison helpers + promote to CPU float64 before measuring. + """ + aniso = kind == "aniso" + scene = scene_fine.to(device=device, dtype=dtype) + obs = oracle_fine[f"{kind}_obs"].to(device=device, dtype=dtype) + leaves = scene.leaves(aniso=aniso) + F = H.density_to_F(scene, H.splat_direct(scene, kernel, *leaves, aniso=aniso)) + got = torch.autograd.grad(H.ls_target(F, obs), leaves) + ref = oracle_fine[f"{kind}_grads"] + + print(f"\n {device.type}/{dtype}/{kernel}/{kind}") + for name, g, r in zip(_LEAF_NAMES, got, ref): + rel, cos = rel_error(g, r), cosine_similarity(g, r) + print(f" {name:10s} rel {rel:.3e} cos {cos:.8f} ratio {gradnorm_ratio(g, r):.4f}") + assert rel < RTOL_GRADIENT_SYNTHETIC, ( + f"{device.type}/{dtype}/{kernel}/{kind} {name}: rel {rel:.3e}" + ) + assert cos > COS_MIN_GRADIENT_SYNTHETIC, ( + f"{device.type}/{dtype}/{kernel}/{kind} {name}: cos {cos:.6f}" + ) + + +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +@pytest.mark.parametrize("device,dtype,kernel", DEVICE_DTYPE_KERNELS) +def test_kernel_gradients_agree_with_portable(scene_fine, oracle_fine, device, dtype, kernel, kind): + """Kernel-vs-portable gradients on the same device, gated tightly. + + The oracle gate above is loose by necessity -- it absorbs the real discretization + error. This one is tight, because two kernels applying the same truncation contract to + the same inputs differ only in accumulation order. It is what would catch a backward + that is wrong in the same direction as the discretization residual. + """ + if kernel == "portable": + pytest.skip("portable is the reference for this comparison") + aniso = kind == "aniso" + scene = scene_fine.to(device=device, dtype=dtype) + obs = oracle_fine[f"{kind}_obs"].to(device=device, dtype=dtype) + + def grads(name): + leaves = scene.leaves(aniso=aniso) + F = H.density_to_F(scene, H.splat_direct(scene, name, *leaves, aniso=aniso)) + return torch.autograd.grad(H.ls_target(F, obs), leaves) + + got, ref = grads(kernel), grads("portable") + tol = RTOL_BACKEND_GRAD_F32 if dtype is torch.float32 else RTOL_BACKEND_GRAD_F64 + for name, g, r in zip(_LEAF_NAMES, got, ref): + rel = rel_error(g, r) + print(f" {device.type}/{dtype}/{kernel}/{kind} {name:10s} vs portable rel {rel:.3e}") + assert rel < tol, f"{device.type}/{dtype}/{kernel}/{kind} {name}: rel {rel:.3e}" + + +# --------------------------------------------------------------------------- +# 6. Direct-summation kernels vs the eager oracle +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +@pytest.mark.parametrize("device,dtype,kernel", DS_DEVICE_DTYPE_KERNELS) +def test_ds_kernel_gradients_match_eager_oracle(scene_fine, device, dtype, kernel, kind): + """Gradients from one DS kernel against the eager oracle, all three leaves. + + Absorbs ``test_gradient_correctness.py``'s two Triton DS gradient tests, which lived + beside geometry and ADP-restraint tests and compared against ``_eager_*`` at + ``min_cos=0.999, ratio_tol=1e-2`` while never checking forward values. Here the forward + comparison is :func:`test_forward.test_ds_kernel_matches_eager_oracle` and the gate is + near precision, because these are two analytic implementations with no discretization + between them. + + Coordinates are **fractional** -- direct summation's convention, unlike the splat + kernels' Cartesian. + """ + aniso = kind == "aniso" + scene = scene_fine.to(device=device, dtype=dtype) + third_ref = scene_fine.u6 if aniso else scene_fine.adp + + with torch.no_grad(): + obs = H.synthetic_obs(H.ds_direct(scene_fine, "eager", aniso=aniso)) + + def grads(sc, name, third_src, obs_local): + leaves = tuple( + t.clone().requires_grad_(True) + for t in (sc.xyz_frac, sc.occ, third_src) + ) + F = H.ds_direct(sc, name, *leaves, aniso=aniso) + return torch.autograd.grad(H.ls_target(F, obs_local), leaves) + + ref = grads(scene_fine, "eager", third_ref, obs) + got = grads( + scene, kernel, scene.u6 if aniso else scene.adp, obs.to(device=device, dtype=dtype) + ) + + tol = RTOL_DS_F32 if dtype is torch.float32 else RTOL_DS_F64 + print(f"\n {device.type}/{dtype}/{kernel}/{kind}") + for name, g, r in zip(("xyz_frac", "occ", "adp_or_u"), got, ref): + rel, cos = rel_error(g, r), cosine_similarity(g, r) + print(f" {name:10s} rel {rel:.3e} cos {cos:.9f}") + assert rel < tol, f"{device.type}/{dtype}/{kernel}/{kind} {name}: rel {rel:.3e}" + assert cos > COS_MIN_DS, f"{device.type}/{dtype}/{kernel}/{kind} {name}: cos {cos:.7f}" diff --git a/tests/unit/structure_factor/test_second_order.py b/tests/unit/structure_factor/test_second_order.py new file mode 100644 index 00000000..b840f3c8 --- /dev/null +++ b/tests/unit/structure_factor/test_second_order.py @@ -0,0 +1,434 @@ +"""Second derivatives: FD -> direct summation -> the FFT/splat route. + +Why direct summation and not finite differences for the map route +---------------------------------------------------------------- +The two tests this module replaces -- +``test_kernel_fixes.py::test_cpu_plain_scatter_second_derivative_correct`` and +``::test_cpu_cpp_scatter_second_derivative_correct`` -- compared the map route's HVP +against central differences *of the map route*, and were failing. + +The reference was the problem, not the kernel -- but not for the reason an earlier draft +of this module claimed. That draft argued finite differences *diverge* on a truncated +density, because the cull surface moves with the atom and each voxel crossing contributes +a fixed jump. Measured, that does not happen: FD converges normally, reaching rel 5.2e-04 +against the map route's own autograd HVP. At a 3 sigma cutoff the surface density is ~1% +of peak, so the jumps are too small to matter. + +The actual problem is that FD asks the wrong question. It differentiates the *same* +discretized function the kernel does, so it can only confirm that autograd correctly +differentiates whatever the kernel computes. Measured here: the map route's HVP agrees +with itself to 5e-04 while sitting 2.3e-02 from the analytic answer. An FD-based gate at +5e-03 -- what the replaced tests used -- passes comfortably while the derivative under +test is 2.3e-02 wrong. + +Direct summation has no grid and no truncation, so FD *is* a valid reference for it +(measured 1.19e-10), and it is then an independent reference for the map route. +:func:`test_finite_differences_cannot_detect_map_route_error` measures all of this, so the +reasoning is pinned by numbers rather than asserted in prose. + +The regression the original two tests guarded is preserved: a graph-less C++ backward +that produced a measured cosine of ~0.57 while first-order gradients stayed correct. +That is still caught here, against a reference that converges. +""" + +from __future__ import annotations + +import pytest +import torch + +from tests.helpers.grad_asserts import cosine_similarity, hvp, hvp_central_fd, rel_error +from torchref.base.direct_summation.dispatch import _eager_aniso, _eager_iso +from torchref.base.electron_density._backends import DENSITY_BACKENDS + +from . import ( + ATOL_GRADCHECK, + COS_MIN, + EPS_GRADCHECK, + RTOL_DS_HVP_VS_FD, + RTOL_GRADCHECK, + RTOL_HVP, +) +from . import helpers as H +from .conftest import DEVICE_DTYPE_KERNELS + +pytestmark = pytest.mark.unit + + +def _lbl(pin): + """Label for the backend a pin parametrization selected.""" + return "portable" if pin else "default" + +_DTYPES = [torch.float32, torch.float64] # float32 first: production dtype + + +# --------------------------------------------------------------------------- +# 1. The oracle is sound at second order +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +def test_ds_gradgradcheck(scene_small, kind): + """``gradgradcheck`` on the eager DS path -- the first in this repo. + + ``grep gradgradcheck`` over ``torchref/`` and ``tests/`` previously returned nothing, + so no second derivative anywhere had been checked against finite differences at the + autograd level. Everything second-order was either cross-backend parity or an HVP + against an FD reference on the truncated map route. + """ + s = scene_small + if kind == "iso": + fn = lambda x, o, a: _eager_iso(s.hkl, s.s, x, o, a, s.A, s.B, None) # noqa: E731 + third = s.adp + else: + fn = lambda x, o, u: _eager_aniso( # noqa: E731 + s.hkl, s.s_vec, x, o, u, s.A, s.B, None + ) + third = s.u6 + + args = tuple(t.clone().requires_grad_(True) for t in (s.xyz_frac, s.occ, third)) + assert torch.autograd.gradgradcheck( + fn, args, eps=EPS_GRADCHECK, atol=ATOL_GRADCHECK, rtol=RTOL_GRADCHECK + ) + + +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +def test_ds_hvp_matches_finite_differences(scene_small, kind): + """The oracle's HVP against central differences of its own gradient. + + Complements ``gradgradcheck``: that checks the full double-backward machinery + elementwise on a tiny input, this checks the specific contraction the map-route tests + use, so the two links are compared like for like. + """ + s = scene_small + fn = H.ds_aniso_oracle if kind == "aniso" else H.ds_iso_oracle + third = s.u6 if kind == "aniso" else s.adp + with torch.no_grad(): + obs = H.synthetic_obs(fn(s, s.xyz, s.occ, third)) + loss = lambda x: H.ls_target(fn(s, x, s.occ, third), obs) # noqa: E731 + + g = torch.Generator().manual_seed(17) + v = torch.randn(s.xyz.shape, generator=g, dtype=s.xyz.dtype) + v /= v.norm() + + got = hvp(loss, s.xyz, v) + ref = hvp_central_fd(loss, s.xyz, v, eps=1e-6) + rel, cos = rel_error(got, ref), cosine_similarity(got, ref) + print(f"\n {kind}: DS HVP vs central FD -- rel {rel:.3e} cos {cos:.10f}") + assert rel < RTOL_DS_HVP_VS_FD + assert cos > 1.0 - 1e-8 + + +# --------------------------------------------------------------------------- +# 2. The first-order-only contract of the public DS API, made explicit +# --------------------------------------------------------------------------- +def test_public_ds_api_raises_on_double_backward(scene_small): + """``ds_iso`` and ``SfDS`` are first-order only. Assert it, so it is documented. + + ``_CheckpointedSF.backward`` calls ``torch.autograd.grad`` without + ``create_graph=True`` on detached copies + (``torchref/base/direct_summation/dispatch.py:197``), so the returned gradients carry + no graph. It is *not* decorated ``@once_differentiable``, so nothing warns you; the + failure surfaces only when a second derivative is requested. + + It does at least fail loudly rather than returning zeros -- verified: it raises + ``element 0 of tensors does not require grad``. This test pins that, so if anyone + later makes the checkpointed path double-differentiable it fails here and the oracle + guidance in this package's docstring gets revisited rather than quietly going stale. + """ + from torchref.base.direct_summation.dispatch import ds_iso + + s = scene_small + x = s.xyz_frac.clone().requires_grad_(True) + F = ds_iso(s.hkl, s.s, x, s.occ, s.adp, s.A, s.B) + with torch.no_grad(): + obs = H.synthetic_obs(F) + (g1,) = torch.autograd.grad(H.ls_target(F, obs), x, create_graph=True) + + v = torch.ones_like(x) + with pytest.raises(RuntimeError, match="does not require grad"): + torch.autograd.grad((g1 * v).sum(), x) + + +def test_eager_ds_survives_double_backward_where_public_api_does_not(scene_small): + """The positive half of the pair above: the eager path *does* compose. + + Together these two tests state the rule this package depends on -- use ``_eager_*`` + as the second-order oracle, never ``ds_*``/``SfDS`` -- as executable fact rather than + as a comment that can drift. + """ + s = scene_small + x = s.xyz_frac.clone().requires_grad_(True) + F = _eager_iso(s.hkl, s.s, x, s.occ, s.adp, s.A, s.B, None) + with torch.no_grad(): + obs = H.synthetic_obs(F.detach()) + (g1,) = torch.autograd.grad(H.ls_target(F, obs), x, create_graph=True) + (out,) = torch.autograd.grad((g1 * torch.ones_like(x)).sum(), x) + assert torch.isfinite(out).all() and out.abs().sum() > 0 + + +# --------------------------------------------------------------------------- +# 3. DS -> FFT at second order +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("pin", [False, True], ids=["default", "portable"]) +@pytest.mark.parametrize("dtype", _DTYPES, ids=["f32", "f64"]) +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +def test_fft_hvp_matches_ds(scene_fine, oracle_fine, kind, dtype, pin): + """The map route's HVP against the DS oracle's HVP. + + This is the replacement for the two failing FD-based tests, and it covers strictly + more: both dtypes and both backends, where the originals covered float64-plain and + float32-C++ only, each with its own hand-tuned ``eps``. + + Both sides contract with the same seeded direction ``v`` (from the oracle fixture), + since an HVP is only comparable along a shared direction. + """ + aniso = kind == "aniso" + sf = H.sf_fft_for(scene_fine, dtype) + v = oracle_fine[f"{kind}_v"] + _, occ, third = scene_fine.leaves(aniso=aniso, requires_grad=False) + + def loss(x): + return H.ls_target( + H.fft_sf(scene_fine, sf, x, occ, third, aniso=aniso), + oracle_fine[f"{kind}_obs"], + ) + + with H.maybe_portable(pin): + got = hvp(loss, scene_fine.xyz, v) + ref = oracle_fine[f"{kind}_hvp"] + + rel, cos = rel_error(got, ref), cosine_similarity(got, ref) + print(f"\n {kind}/{dtype}/{_lbl(pin)}: HVP vs DS -- rel {rel:.3e} cos {cos:.8f}") + assert cos > COS_MIN, f"{kind}/{dtype}/{_lbl(pin)}: HVP direction, cos {cos:.6f}" + assert rel < RTOL_HVP, f"{kind}/{dtype}/{_lbl(pin)}: HVP magnitude, rel {rel:.3e}" + + +@pytest.mark.parametrize("dtype", _DTYPES, ids=["f32", "f64"]) +def test_fused_cpu_kernel_uses_the_double_backward_fallback( + scene_fine, oracle_fine, dtype, monkeypatch +): + """Non-vacuity guard for the AUTO second-order path. + + The fused C++ sphere splat has a hand-written first-order backward with no graph, so + under ``create_graph=True`` it routes through ``_double_backward_vjp``, which re-runs + the forward through the portable torch splat on the *saved* leaves and differentiates + that. Without this guard, ``test_fft_hvp_matches_ds[auto]`` would still pass if AUTO + silently fell back to the eager path for the whole forward -- and would then be + testing EAGER twice under two names. + + Asserts the fallback is entered, which is only meaningful alongside a separate + assertion that the fused kernel was used in the first place. + """ + from torchref.base.electron_density.kernels.cpu import sphere_splat + + if not sphere_splat.sphere_splat_available(): + pytest.skip(f"fused CPU sphere splat unavailable: {sphere_splat.last_error()}") + + calls = {"fallback": 0} + original = sphere_splat._double_backward_vjp + + def counting(*args, **kwargs): + calls["fallback"] += 1 + return original(*args, **kwargs) + + monkeypatch.setattr(sphere_splat, "_double_backward_vjp", counting) + + sf = H.sf_fft_for(scene_fine, dtype) + v = torch.ones_like(scene_fine.xyz) + _, occ, adp = scene_fine.leaves(requires_grad=False) + obs = oracle_fine["iso_obs"] + + def loss(x): + return H.ls_target(H.fft_sf(scene_fine, sf, x, occ, adp), obs) + + with H.maybe_portable(False): + out = hvp(loss, scene_fine.xyz, v) + + assert calls["fallback"] > 0, ( + "AUTO produced an HVP without entering _double_backward_vjp, so the fused " + "kernel was not on the path and this dtype's second-order coverage is vacuous" + ) + assert torch.isfinite(out).all() and out.abs().sum() > 0 + + +# --------------------------------------------------------------------------- +# 4. The measurement that justifies choosing DS over finite differences +# --------------------------------------------------------------------------- +def test_finite_differences_cannot_detect_map_route_error(scene_fine, oracle_fine): + """Finite differences confirm self-consistency, not correctness. Measured. + + This test was originally written to demonstrate an *FD plateau* on the map route -- + the argument being that a hard truncation whose cull surface moves with the atom + contributes a fixed jump per crossing, so the spurious term shrinks no faster than the + divisor. That argument is inherited from an earlier draft and the measurement does not + support it: FD converges normally here, reaching rel 5.2e-04 against the map route's + own autograd HVP at ``eps = 1e-2``, with the usual V-shape (truncation-dominated at + large ``eps``, round-off at small). At the 3 sigma cutoff the density at the surface is + ~1% of peak, so the jumps are evidently too small to matter. + + The measurement supports a stronger conclusion instead. The map route's HVP agrees + with *itself* to 5e-04 while sitting 2.3e-02 away from the analytic answer. So an + FD-based test gated at the 5e-03 the replaced tests used passes comfortably while the + derivative under test is 2.3e-02 wrong -- FD cannot see that error at all, because it + differentiates the same discretized function. Only an independent reference can. + + That is why this package uses DS. Not because FD diverges, but because FD is measuring + the wrong thing: it validates that autograd correctly differentiates whatever the + kernel computes, and is blind by construction to whether the kernel computes the right + thing. + """ + sf = H.sf_fft_for(scene_fine) + v = oracle_fine["iso_v"] + _, occ, adp = scene_fine.leaves(requires_grad=False) + obs = oracle_fine["iso_obs"] + + def loss(x): + return H.ls_target(H.fft_sf(scene_fine, sf, x, occ, adp), obs) + + auto = hvp(loss, scene_fine.xyz, v) + ref_ds = oracle_fine["iso_hvp"] + + print("\n eps rel vs autograd(map) rel vs DS oracle") + rows = [] + for eps in (1e-4, 1e-3, 1e-2, 1e-1): + fd = hvp_central_fd(loss, scene_fine.xyz, v, eps=eps) + r_self, r_ds = rel_error(fd, auto), rel_error(fd, ref_ds) + rows.append((eps, r_self, r_ds)) + print(f" {eps:<9.0e} {r_self:18.3e} {r_ds:16.3e}") + + print(f" autograd(map) vs DS oracle: {rel_error(auto, ref_ds):.3e}") + + best_self = min(r for _, r, _ in rows) + best_vs_ds = min(r for _, _, r in rows) + map_vs_ds = rel_error(auto, ref_ds) + + # 1. FD does converge against the map route's own autograd. + assert best_self < 5e-3, ( + f"FD only reached rel {best_self:.2e} against the map route's autograd HVP. This " + "test's premise is that it converges; if that has stopped being true the " + "docstring's reasoning needs rewriting, not the tolerance." + ) + # 2. And is nonetheless blind to the real error, which is what matters. + assert best_vs_ds > 4 * best_self, ( + f"FD sits {best_vs_ds:.2e} from the analytic answer but {best_self:.2e} from the " + "map route's autograd. Those being comparable would mean the map route is now " + "accurate enough that FD and DS are interchangeable references, and this test " + "no longer demonstrates anything -- re-examine rather than retune." + ) + assert map_vs_ds > 5e-3, ( + f"the map route's HVP is now within {map_vs_ds:.2e} of the analytic answer, " + "which is inside the 5e-3 gate the replaced FD-based tests used -- so those " + "tests would no longer have been misleading and this rationale is stale" + ) + + +def test_fft_hvp_matches_ds_real_structure(gemmi_aniso_grad, oracle_aniso_grad): + """Second derivatives on 7L84 at production sampling, in the production dtype. + + The synthetic HVP sweep above covers dtype and backend combinations; this is the one + that speaks to production, for the same reason as its first-order counterpart in + ``test_gradients.py``. + """ + scene, _ = gemmi_aniso_grad + sf = H.sf_fft_for(scene, torch.float32) + v = oracle_aniso_grad["aniso_v"] + _, occ, u6 = scene.leaves(aniso=True, requires_grad=False) + obs = oracle_aniso_grad["aniso_obs"] + + def loss(x): + return H.ls_target(H.fft_sf(scene, sf, x, occ, u6, aniso=True), obs) + + got = hvp(loss, scene.xyz, v) + ref = oracle_aniso_grad["aniso_hvp"] + rel, cos = rel_error(got, ref), cosine_similarity(got, ref) + print(f"\n7L84 P1 HVP vs DS: rel {rel:.3e} cos {cos:.8f}") + assert cos > COS_MIN, f"7L84 HVP direction: cos {cos:.6f}" + assert rel < RTOL_HVP, f"7L84 HVP magnitude: rel {rel:.3e}" + + + +# --------------------------------------------------------------------------- +# 5. Every production kernel, on every device it ships on +# --------------------------------------------------------------------------- +# Second order is where the kernels genuinely differ, so this section is not simply the +# gradient section with another derivative. Only two of the four families are +# double-differentiable: +# +# add_*_cpu_sphere_var yes -- via ``_double_backward_vjp``, which re-derives the VJP +# through the portable splat on the saved leaves +# add_*_plain_var yes -- pure autograd over ``scatter_add`` +# add_*_mps_var no -- ``mps/variable_radius.py:12``: "Backward is first-order +# only (like CUDA); double backward needs the portable splat" +# WorkQueueGridDensity* no -- same, no ``create_graph`` path in the Triton backward +# +# So an accelerator HVP cannot be compared against the oracle: there is no HVP to compare. +# What is testable, and tested below, is that those kernels **raise** rather than returning +# a silently wrong (or zero) second derivative -- the same failure mode this repo already +# fixed once for the C++ scatter, where a graph-less backward gave cosine 0.57 while first +# derivatives stayed correct. + +#: Read from the production table rather than restated. This was a hand-written set, i.e. a +#: third copy of the same fact (the other two being the kernel docstrings and the table), +#: and the kind that goes stale silently: a kernel gaining a ``create_graph`` path would +#: leave its HVP untested, and one losing it would fail confusingly. +_DOUBLE_DIFFERENTIABLE = { + b.name for b in DENSITY_BACKENDS.backends if b.second_order +} + + +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +@pytest.mark.parametrize("device,dtype,kernel", DEVICE_DTYPE_KERNELS) +def test_kernel_hvp_matches_ds(scene_fine, oracle_fine, device, dtype, kernel, kind): + """HVP of a double-differentiable kernel against the oracle's, on every device. + + Covers the portable splat running *on an accelerator*, which nothing tested before: + it is what ``force_portable`` and the CUDA/MPS float64 fallthrough actually execute, and + it is the path a Hessian-based optimizer lands on there. + """ + if kernel not in _DOUBLE_DIFFERENTIABLE: + pytest.skip(f"{kernel} is first-order only; see test_kernel_rejects_double_backward") + aniso = kind == "aniso" + scene = scene_fine.to(device=device, dtype=dtype) + obs = oracle_fine[f"{kind}_obs"].to(device=device, dtype=dtype) + v = oracle_fine[f"{kind}_v"].to(device=device, dtype=dtype) + _, occ, third = scene.leaves(aniso=aniso, requires_grad=False) + + def loss(x): + return H.ls_target( + H.density_to_F(scene, H.splat_direct(scene, kernel, x, occ, third, aniso=aniso)), + obs, + ) + + got = hvp(loss, scene.xyz, v) + ref = oracle_fine[f"{kind}_hvp"] + rel, cos = rel_error(got, ref), cosine_similarity(got, ref) + print(f"\n {device.type}/{dtype}/{kernel}/{kind}: HVP rel {rel:.3e} cos {cos:.8f}") + assert cos > COS_MIN, f"{device.type}/{dtype}/{kernel}/{kind}: HVP cos {cos:.6f}" + assert rel < RTOL_HVP, f"{device.type}/{dtype}/{kernel}/{kind}: HVP rel {rel:.3e}" + + +@pytest.mark.parametrize("kind", ["iso", "aniso"]) +@pytest.mark.parametrize("device,dtype,kernel", DEVICE_DTYPE_KERNELS) +def test_kernel_rejects_double_backward(scene_fine, oracle_fine, device, dtype, kernel, kind): + """First-order-only kernels must raise under ``create_graph=True``, not return garbage. + + The accelerator kernels have hand-written backwards that return tensors carrying no + graph, and neither is decorated ``@once_differentiable`` -- so nothing warns, and the + only signal is whatever autograd does when asked for a second derivative. Pinning that + it *raises* is what distinguishes "unsupported" from "silently wrong". + + The positive half of the pair is :func:`test_kernel_hvp_matches_ds`; together they say + which kernels an optimizer may take a Hessian through. + """ + if kernel in _DOUBLE_DIFFERENTIABLE: + pytest.skip(f"{kernel} supports double backward; see test_kernel_hvp_matches_ds") + aniso = kind == "aniso" + scene = scene_fine.to(device=device, dtype=dtype) + obs = oracle_fine[f"{kind}_obs"].to(device=device, dtype=dtype) + _, occ, third = scene.leaves(aniso=aniso, requires_grad=False) + + x = scene.xyz.clone().requires_grad_(True) + F = H.density_to_F(scene, H.splat_direct(scene, kernel, x, occ, third, aniso=aniso)) + (g1,) = torch.autograd.grad(H.ls_target(F, obs), x, create_graph=True) + + with pytest.raises(RuntimeError) as exc: + torch.autograd.grad((g1 * torch.ones_like(x)).sum(), x) + print(f"\n {device.type}/{dtype}/{kernel}/{kind} raised: {str(exc.value)[:70]}") diff --git a/tests/unit/symmetry/test_hkl_symmetry_gemmi.py b/tests/unit/symmetry/test_hkl_symmetry_gemmi.py new file mode 100644 index 00000000..8953b20e --- /dev/null +++ b/tests/unit/symmetry/test_hkl_symmetry_gemmi.py @@ -0,0 +1,158 @@ +"""Regression tests for reciprocal-space (Miller-index) symmetry against gemmi. + +These guard the ``h' = h·R = Rᵀ·h`` reciprocal-space convention used by +``SpaceGroup.apply_to_hkl``. A previous bug applied the real-space transform +``R·h`` instead, which is only correct when the fractional rotation matrix is +symmetric. For space groups whose rotation matrices are non-symmetric +(trigonal, hexagonal, and permutation-type cubic operations) ``R·h`` produced +the wrong set of symmetry equivalents, corrupting the centric flags +(``is_centric_from_hkl``) and epsilon multiplicities (``epsilon_from_hkl``) +that feed French-Wilson intensity conversion and ML sigma_A weighting. + +Ground truth comes from gemmi: +- transform: ``gemmi.Op.apply_to_hkl`` +- centric: ``GroupOps.is_reflection_centric`` +- epsilon: ``GroupOps.epsilon_factor_without_centering`` (Friedel-doubled + for centric reflections, matching the Friedel-aware count in + ``epsilon_from_hkl``) +""" + +import numpy as np +import pytest +import torch + +from torchref.config import get_default_device, get_float_dtype, get_int_dtype + +# Space groups whose rotation matrices are non-symmetric — these are the ones +# that regress if apply_to_hkl uses R·h instead of Rᵀ·h. Plus orthorhombic / +# tetragonal controls (symmetric matrices) that must remain correct either way. +_TRIGONAL_HEXAGONAL = ["P 32 2 1", "P 31 2 1", "P 61 2 2", "P 6 2 2", "P 3 1 2"] +_CONTROLS = ["P 21 21 21", "P 43 21 2", "P 1"] +_SPACE_GROUPS = _TRIGONAL_HEXAGONAL + _CONTROLS + +# A spread of reflections including negative indices and axial/general cases. +_HKLS = [ + (1, 0, 0), (0, 1, 0), (0, 0, 1), (1, 1, 0), (2, 1, 0), (1, -1, 0), + (1, 2, 3), (0, 1, 2), (3, 1, 0), (0, 0, 2), (1, 0, 1), (-1, 2, -3), + (2, 2, 1), (1, 1, 1), +] + + +def _hkl_tensor(): + """Miller indices on the configured default device and integer dtype. + + Kept at the config defaults rather than a hardcoded float64/CPU so the test + exercises the same precision and device production does, and stays + MPS-compatible (MPS has no float64). The arithmetic here is exact in + float32 regardless: fractional rotation matrices are integer-valued + (gemmi's ``op.rot / 24``) and the Miller indices are small, so every + product and sum is well inside the exactly-representable integer range. + """ + return torch.tensor(_HKLS, dtype=get_int_dtype(), device=get_default_device()) + + +@pytest.mark.unit +@pytest.mark.parametrize("sg_name", _SPACE_GROUPS) +def test_apply_to_hkl_matches_gemmi_transform(sg_name): + """apply_to_hkl must reproduce gemmi's per-operation reciprocal transform. + + Compared as the *set* of equivalents per reflection so the test is + insensitive to operation ordering between torchref and gemmi. + """ + import gemmi + + from torchref.symmetry import SpaceGroup + + sg = SpaceGroup(sg_name) + hkl = _hkl_tensor() + + # torchref: (N, 3, ops) -> per-reflection set of equivalents + out = sg.apply_to_hkl(hkl.to(get_float_dtype())) + out_int = torch.round(out).to(torch.int64).cpu().numpy() # (N, 3, ops) + + gemmi_ops = list(gemmi.SpaceGroup(sg_name).operations().sym_ops) + + for n, h in enumerate(_HKLS): + got = {tuple(out_int[n, :, o]) for o in range(out_int.shape[2])} + expected = {tuple(op.apply_to_hkl(h)) for op in gemmi_ops} + assert got == expected, ( + f"{sg_name} hkl={h}: apply_to_hkl equivalents {sorted(got)} " + f"!= gemmi {sorted(expected)}" + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("sg_name", _SPACE_GROUPS) +def test_is_centric_from_hkl_matches_gemmi(sg_name): + """Centric flags must match gemmi's is_reflection_centric.""" + import gemmi + + from torchref.base.french_wilson import is_centric_from_hkl + + ops = gemmi.SpaceGroup(sg_name).operations() + hkl = _hkl_tensor() + + got = is_centric_from_hkl(hkl, sg_name).cpu().numpy().astype(bool) + expected = np.array([ops.is_reflection_centric(h) for h in _HKLS], dtype=bool) + + assert np.array_equal(got, expected), ( + f"{sg_name}: centric {got.astype(int)} != gemmi {expected.astype(int)}" + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("sg_name", _SPACE_GROUPS) +def test_epsilon_from_hkl_matches_gemmi(sg_name): + """Friedel-aware epsilon must match gemmi's epsilon (doubled for centrics).""" + import gemmi + + from torchref.base.targets.xray_ml_sigmaa import epsilon_from_hkl + from torchref.symmetry import SpaceGroup + + ops = gemmi.SpaceGroup(sg_name).operations() + sg = SpaceGroup(sg_name) + hkl = _hkl_tensor() + + got = epsilon_from_hkl(hkl, sg).cpu().numpy().round().astype(int) + expected = np.array( + [ + ops.epsilon_factor_without_centering(h) + * (2 if ops.is_reflection_centric(h) else 1) + for h in _HKLS + ], + dtype=int, + ) + + assert np.array_equal(got, expected), ( + f"{sg_name}: epsilon {got} != gemmi-derived {expected}" + ) + + +@pytest.mark.unit +def test_apply_to_hkl_is_transpose_not_plain_rotation(): + """Explicit guard on the exact bug: apply_to_hkl == Rᵀ·h, not R·h. + + Uses P3, whose 3-fold rotation matrix is non-symmetric, so R·h and Rᵀ·h + give genuinely different results. + """ + from torchref.symmetry import SpaceGroup + + sg = SpaceGroup("P 3") + # Use the SpaceGroup's own dtype (config default) -- no hardcoded float64, + # so this runs on MPS as well as CPU/CUDA. + mats = sg.matrices # (ops, 3, 3) + hkl = torch.tensor([[1, 2, 3]], dtype=mats.dtype, device=mats.device) + + out = sg.apply_to_hkl(hkl) # (1, 3, ops) + + # Correct reciprocal transform: Rᵀ·h + expected_t = torch.einsum("oji,nj->nio", mats, hkl) + # The old (buggy) real-space transform: R·h + plain_r = torch.einsum("oij,nj->nio", mats, hkl) + + assert torch.allclose(out, expected_t), "apply_to_hkl should compute Rᵀ·h" + # For P3 the two conventions must actually differ (non-symmetric matrix), + # otherwise this test would not detect a regression. + assert not torch.allclose(expected_t, plain_r), ( + "P3 rotation matrices should be non-symmetric; test cannot detect the bug" + ) diff --git a/tests/unit/symmetry/test_phase_convention.py b/tests/unit/symmetry/test_phase_convention.py new file mode 100644 index 00000000..2d6436e4 --- /dev/null +++ b/tests/unit/symmetry/test_phase_convention.py @@ -0,0 +1,233 @@ +"""Phase-shift sign convention for reciprocal-space symmetry operations. + +These tests pin the phase contract of :func:`expand_hkl`, :func:`canonicalize_hkl` +and :func:`reduce_hkl` against an *independent* reference: a direct structure-factor +summation over an explicitly symmetry-expanded atom set. That is the definition of a +structure factor, so it cannot drift in step with the implementation the way a +path-vs-path comparison can. + +Why this file exists +-------------------- +All three functions computed ``+2π h·t`` where the correct shift is ``-2π h·t`` +(and, for the Friedel-flipped rows of ``canonicalize_hkl``, ``+2π h·t`` -- the sign +is *not* uniform there; see that function's comment). The bug survived because the +residual error of the wrong sign is ``4π h·t mod 2π``, which is **exactly zero** for +2₁ screw axes and for centring translations -- and the pre-existing phase test +(``test_canonicalize_hkl.py::test_phase_roundtrip``) parametrised only over +``P21, P212121, C2, P4``, every one of which falls in that blind spot. + +The parametrisation below deliberately spans all three regimes: + +=================================== =========================== +translation error if the sign is flipped +=================================== =========================== +2₁ screws, centring (P21, C2, ...) 0 (invisible) +4₁/4₃ screws (P41, P43, P43212) π +3₁/6₁ screws (P31, P61, P3121, ...) 2π/3 +=================================== =========================== + +Do not narrow this list. ``test_wrong_sign_would_be_detected`` guards against it +becoming vacuous. +""" + +import numpy as np +import pytest +import torch + +from torchref.config import get_float_dtype +from torchref.symmetry.reciprocal_symmetry import ( + canonicalize_hkl, + expand_hkl, + reduce_hkl, +) +from torchref.symmetry.spacegroup import SpaceGroup + +# Groups spanning the three regimes above. P1 is the degenerate control (no +# translations at all); the screw-axis groups are the ones with real signal. +SPACE_GROUPS = [ + "P 1", + "P 21", + "P 21 21 21", + "C 1 2 1", + "P 41", + "P 43", + "P 31", + "P 61", + "P 43 21 2", + "P 31 2 1", + "P 65 2 2", + "I 41", +] + +# Groups where a flipped sign produces a non-zero phase error. Kept separate so the +# anti-vacuity test can assert the suite is actually sensitive to the regression. +SENSITIVE_GROUPS = ["P 41", "P 43", "P 31", "P 61", "P 43 21 2", "P 31 2 1", "P 65 2 2"] + +# Comfortably tighter than the smallest real error (2π/3 ≈ 2.09 rad) yet loose +# enough for the float32 default dtype the functions compute in (~1e-6 observed). +ATOL_RAD = 1e-4 + + +def _wrap(a): + """Wrap angles into (-π, π] so 2π-equivalent phases compare equal.""" + return (np.asarray(a) + np.pi) % (2 * np.pi) - np.pi + + +def _reference_model(spacegroup, n_atoms=10, seed=1): + """A symmetry-obeying atom set plus an exact ``F(h)`` evaluator. + + The atom list is the asymmetric unit expanded by every symmetry operation, so + the resulting structure genuinely has the space-group symmetry and its + structure factors must obey ``arg F(hR) - arg F(h) = -2π h·t``. Evaluated in + float64 regardless of the library's configured dtype -- this is the reference. + """ + rng = np.random.default_rng(seed) + sg = SpaceGroup(spacegroup, dtype=get_float_dtype(), device=torch.device("cpu")) + rot = sg.matrices.cpu().numpy().astype(np.float64) + trans = sg.translations.cpu().numpy().astype(np.float64) + + xyz_asu = rng.random((n_atoms, 3)) + f_asu = rng.uniform(4.0, 10.0, n_atoms) + xyz = np.concatenate([xyz_asu @ rot[i].T + trans[i] for i in range(len(rot))]) + f = np.concatenate([f_asu] * len(rot)) + + def structure_factor(hkl): + hkl = np.atleast_2d(np.asarray(hkl, dtype=np.float64)) + return (f[None, :] * np.exp(2j * np.pi * (hkl @ xyz.T))).sum(axis=1) + + return structure_factor + + +def _random_hkl(n=8, seed=3): + """Random non-zero Miller indices.""" + rng = np.random.default_rng(seed) + out = [] + while len(out) < n: + h = tuple(int(x) for x in rng.integers(-6, 7, 3)) + if h != (0, 0, 0): + out.append(h) + return torch.tensor(np.array(out), dtype=torch.int32) + + +@pytest.mark.unit +@pytest.mark.parametrize("spacegroup", SPACE_GROUPS) +def test_expand_hkl_phase_contract(spacegroup): + """``phi_expanded = phi_orig[indices] + phase_shifts`` must be the true phase. + + This is the contract stated in :func:`expand_hkl`'s own docstring. Checked with + ``include_friedel=False``: the Friedel half cannot be expressed as an additive + shift at all (``phi(-h) = -phi(h)`` is a conjugation, not an offset), so the + documented contract only applies to the pure-rotation expansion. + """ + fcalc = _reference_model(spacegroup) + hkl = _random_hkl() + phi_in = np.angle(fcalc(hkl.numpy())) + + hkl_p1, indices, shifts = expand_hkl( + hkl, spacegroup, include_friedel=False, remove_absences=True + ) + got = phi_in[indices.cpu().numpy()] + shifts.cpu().numpy() + expected = np.angle(fcalc(hkl_p1.numpy())) + + err = np.abs(_wrap(got - expected)).max() + assert err < ATOL_RAD, ( + f"{spacegroup}: expand_hkl phase contract violated by {err:.4f} rad. " + f"Expected arg F(hR) - arg F(h) == -2*pi*h.t." + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("spacegroup", SPACE_GROUPS) +def test_canonicalize_hkl_phase_contract(spacegroup): + """``phi_canonical = where(friedel, -phi_in, phi_in) + phase_shifts``. + + This is the form :meth:`ReflectionData._canonicalize_in_place` applies. Note the + sign of the shift differs between the Friedel and non-Friedel halves, so a + uniform sign is necessarily wrong for one of them. + """ + fcalc = _reference_model(spacegroup) + # Feed a full-sphere, deliberately non-canonical set so real work is required. + hkl_all, _, _ = expand_hkl( + _random_hkl(), spacegroup, include_friedel=True, remove_absences=True + ) + phi_in = np.angle(fcalc(hkl_all.numpy())) + + canonical, shifts, friedel, sort_idx = canonicalize_hkl( + hkl_all, spacegroup, include_friedel=True + ) + f = friedel.cpu().numpy() + base = np.where(f, -phi_in[sort_idx.cpu().numpy()], phi_in[sort_idx.cpu().numpy()]) + got = base + shifts.cpu().numpy() + expected = np.angle(fcalc(canonical.numpy())) + + err = np.abs(_wrap(got - expected)).max() + assert err < ATOL_RAD, ( + f"{spacegroup}: canonicalize_hkl phase contract violated by {err:.4f} rad " + f"({f.sum()} of {len(f)} rows were Friedel-flipped)." + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("spacegroup", SPACE_GROUPS) +def test_reduce_hkl_phase_contract(spacegroup): + """Every equivalent must reconstruct the ASU phase, not just the first. + + ``reduce_hkl`` returns one row per ASU reflection and a column per equivalent. + Column 0 is typically the identity operation (zero shift), so checking only that + column would pass regardless of the sign -- iterate over all of them. + """ + fcalc = _reference_model(spacegroup) + hkl_p1, _, _ = expand_hkl( + _random_hkl(), spacegroup, include_friedel=False, remove_absences=True + ) + phi_p1 = np.angle(fcalc(hkl_p1.numpy())) + + hkl_asu, red_idx, red_shift = reduce_hkl( + hkl_p1, spacegroup, include_friedel=False + ) + phi_asu = np.angle(fcalc(hkl_asu.numpy())) + idx, sh = red_idx.cpu().numpy(), red_shift.cpu().numpy() + + errs, n_checked = [], 0 + for row in range(idx.shape[0]): + for col in range(idx.shape[1]): + src = idx[row, col] + if src < 0: + continue + n_checked += 1 + errs.append(_wrap(phi_p1[src] + sh[row, col] - phi_asu[row])) + + assert n_checked >= idx.shape[0], "no equivalents were checked" + err = np.abs(errs).max() + assert err < ATOL_RAD, ( + f"{spacegroup}: reduce_hkl phase contract violated by {err:.4f} rad " + f"over {n_checked} equivalents." + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("spacegroup", SENSITIVE_GROUPS) +def test_wrong_sign_would_be_detected(spacegroup): + """Anti-vacuity guard: the tests above must be *able* to fail. + + Re-runs the ``expand_hkl`` check with the shift negated and asserts it breaks. If + someone narrows ``SPACE_GROUPS`` back to only 2₁/centring groups, or a refactor + makes the shifts identically zero, the contract tests would pass trivially and + this test is what notices. + """ + fcalc = _reference_model(spacegroup) + hkl = _random_hkl() + phi_in = np.angle(fcalc(hkl.numpy())) + + hkl_p1, indices, shifts = expand_hkl( + hkl, spacegroup, include_friedel=False, remove_absences=True + ) + expected = np.angle(fcalc(hkl_p1.numpy())) + flipped = phi_in[indices.cpu().numpy()] - shifts.cpu().numpy() + + err = np.abs(_wrap(flipped - expected)).max() + assert err > 0.5, ( + f"{spacegroup}: negating the phase shift changed nothing (max err " + f"{err:.2e} rad), so this space group cannot detect a sign regression " + f"and does not belong in SENSITIVE_GROUPS." + ) diff --git a/tests/unit/symmetry/test_symmetry.py b/tests/unit/symmetry/test_symmetry.py index d04044e1..1ddab2bf 100644 --- a/tests/unit/symmetry/test_symmetry.py +++ b/tests/unit/symmetry/test_symmetry.py @@ -212,8 +212,8 @@ def test_spacegroup_gpu(self, gpu_device): sg = SpaceGroup("P21", device=gpu_device) - assert sg.matrices.device.type == 'cuda' - assert sg.translations.device.type == 'cuda' + assert sg.matrices.device.type == gpu_device.type + assert sg.translations.device.type == gpu_device.type @pytest.mark.unit def test_spacegroup_dtype(self): diff --git a/tests/unit/test_device_conformance.py b/tests/unit/test_device_conformance.py new file mode 100644 index 00000000..ca9ac4ba --- /dev/null +++ b/tests/unit/test_device_conformance.py @@ -0,0 +1,255 @@ +"""Device conformance: constructors place tensors where they were told, and +``.to()`` moves *everything* -- tensors and trackers alike. + +The package resolves a single default device at import and expects every object +to live on it, with ``self.device`` as the authority for where new tensors are +allocated (200+ call sites pass ``device=self.device``). Two failure modes +follow, and both used to be untested: + +* a constructor that builds a sub-object on the *global* default rather than + the device it was handed, producing a split object that only errors much + later, deep inside some unrelated op; +* a ``.to()`` that moves the tensors but leaves ``self.device`` stale, so the + object keeps allocating on the device it just left. + +Before this file, the only test that exercised a real device transition was +CUDA-gated -- meaning zero coverage on CPU-only CI and on Apple silicon. +""" + +from pathlib import Path + +import pytest +import torch + +from torchref.config import canonical_device + +# Absolute ``tests.`` imports, not bare ``helpers.``: the repo carries two +# pytest configs -- ``tests/pytest.ini`` (rootdir ``tests/``, ``pythonpath = ..``) +# and ``[tool.pytest.ini_options]`` in ``pyproject.toml`` (rootdir the repo +# root). Only the repo root is on ``sys.path`` under both, so a bare +# ``helpers.`` import works from ``tests/`` and fails from the repo root. +from tests.helpers.device_asserts import assert_device_consistent, collect_device_map +from tests.helpers.device_cases import ( + CASES, + TARGET_CASES, + UNCOVERED, + UNCOVERED_PREFIXES, +) +from tests.helpers.device_inventory import device_mixin_classes + +_IDS = [c.name for c in CASES] + + +@pytest.mark.unit +@pytest.mark.parametrize("case", CASES, ids=_IDS) +def test_constructed_on_requested_device(case, any_device): + """Everything a constructor builds lands on the device it was given.""" + obj = case.build(any_device) + assert_device_consistent(obj, any_device, name=case.name, ignore=case.ignore) + + +@pytest.mark.unit +@pytest.mark.parametrize("case", CASES, ids=_IDS) +def test_device_round_trip(case, any_device): + """cpu -> device -> cpu, with tensors and trackers following every leg.""" + obj = case.build(torch.device("cpu")) + assert_device_consistent(obj, "cpu", name=case.name, ignore=case.ignore) + + obj.to(any_device) + assert_device_consistent(obj, any_device, name=case.name, ignore=case.ignore) + + obj.to("cpu") + assert_device_consistent(obj, "cpu", name=case.name, ignore=case.ignore) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "case", [c for c in CASES if not c.tensor_free], ids=[c.name for c in CASES if not c.tensor_free] +) +def test_tracker_agrees_with_owned_tensor(case, any_device): + """``obj.device`` must equal a real tensor's device, index and all. + + ``torch.device('mps') != torch.device('mps:0')``, so a tracker carrying the + un-indexed spelling compares unequal to every tensor the object owns even + though nothing is actually misplaced. Callers do write + ``if t.device != self.device``, so the two spellings have to agree. + """ + obj = case.build(any_device) + tensors, _, _, _ = collect_device_map(obj, case.name) + if not tensors: + pytest.skip(f"{case.name} owns no tensors on this path") + tracker = getattr(obj, "device", None) + assert isinstance(tracker, torch.device), f"{case.name}.device is {tracker!r}" + for path, dev in tensors.items(): + assert tracker == dev, ( + f"{case.name}.device == {tracker!r} but {path} is on {dev!r}; " + "these must compare equal, not merely name the same backend" + ) + + +@pytest.mark.unit +def test_every_device_bearing_class_is_accounted_for(): + """A new device-bearing class must be given a case or an explicit excuse. + + Uses a static AST inventory rather than ``DeviceMixin.__subclasses__()``: + the runtime hook only sees classes whose module happens to be imported, so + it would silently under-report exactly when coverage regressed. + """ + package_root = Path(__file__).resolve().parents[2] / "torchref" + found = device_mixin_classes(package_root) + + covered = ( + {c.cls_name for c in CASES} + | {c.cls_name for c in TARGET_CASES} + | set(UNCOVERED) + ) + missing = { + name: where + for name, where in found.items() + if name not in covered + and not any(where.startswith(p) for p in UNCOVERED_PREFIXES) + } + + assert not missing, ( + "device-bearing classes with no conformance case and no entry in " + "helpers.device_cases.UNCOVERED:\n " + + "\n ".join(f"{n} ({w})" for n, w in sorted(missing.items())) + ) + + +@pytest.mark.unit +def test_uncovered_entries_still_exist(): + """Stop ``UNCOVERED`` accumulating excuses for classes that are long gone.""" + package_root = Path(__file__).resolve().parents[2] / "torchref" + found = set(device_mixin_classes(package_root)) + stale = sorted(set(UNCOVERED) - found) + assert not stale, f"UNCOVERED lists classes that no longer exist: {stale}" + + +# --------------------------------------------------------------------------- +# Refinement targets +# +# Targets were the largest hole in the device story: none of them carried a +# ``device`` tracker at all, so ``DeviceMixin``'s machinery was a no-op for +# every one, and their scalar tunables were allocated on CPU in float32 +# regardless of the model's device or the configured dtype. +# --------------------------------------------------------------------------- + +_TARGET_IDS = [c.name for c in TARGET_CASES] + + +@pytest.mark.unit +@pytest.mark.parametrize("case", TARGET_CASES, ids=_TARGET_IDS) +def test_target_follows_its_model_device(case, device_model_bundle, any_device): + """A target's own buffers land beside the model it acts on.""" + target = case.build(device_model_bundle, any_device) + assert_device_consistent(target, any_device, name=case.name, ignore=case.ignore) + + +@pytest.mark.unit +@pytest.mark.parametrize("case", TARGET_CASES, ids=_TARGET_IDS) +def test_target_tracker_is_set(case, device_model_bundle, any_device): + """``target.device`` exists and is truthful. + + Tensor-free targets (``BondTarget``) can only be checked this way -- and + they are exactly the ones that used to have no tracker at all. + """ + target = case.build(device_model_bundle, any_device) + assert isinstance(target.device, torch.device) + assert canonical_device(target.device) == canonical_device(any_device) + + +@pytest.mark.unit +@pytest.mark.parametrize("case", TARGET_CASES, ids=_TARGET_IDS) +def test_target_forward_does_not_reallocate_buffers(case, device_model_bundle, any_device): + """``forward()`` must not move or re-register buffers. + + Three targets used to repair their CPU-resident buffers lazily on the first + forward. That is what this pins shut: allocation inside forward breaks + CUDA-graph capture, and it silently costs a transfer on the hot path. The + old code fails this on the first call and passes on the second. + """ + target = case.build(device_model_bundle, any_device) + + def snapshot(): + return { + n: (b.data_ptr(), b.device, b.dtype) + for n, b in target.named_buffers(recurse=False) + if b is not None + } + + try: + target() + except Exception as exc: # pragma: no cover - target needs state we lack + pytest.skip(f"{case.name} cannot run forward() here: {exc}") + before = snapshot() + target() + assert snapshot() == before, f"{case.name} mutated its buffers during forward()" + + +@pytest.mark.unit +@pytest.mark.parametrize("case", TARGET_CASES, ids=_TARGET_IDS) +def test_empty_target_state_dict_round_trip(case): + """An empty target must round-trip its own state dict strictly. + + The empty-init path exists precisely so ``load_state_dict`` has something + to load into. ``CoordinateSimilarityTarget`` used to fail this: its index + buffers were registered only inside ``_build_atom_map``, which never ran + without models. + """ + cls = type(case.build({"model": None}, torch.device("cpu"))) + shell = cls() + other = cls() + other.load_state_dict(shell.state_dict(), strict=True) + + +@pytest.mark.unit +def test_target_scalar_buffers_follow_config_dtype(pdb_dir, monkeypatch): + """Scalar tunables honour ``dtypes.float``, not a hardcoded float32. + + CPU-only: MPS has no float64. Before the migration every one of these was + ``torch.tensor(float(x))`` -- silently float32 under a float64 config. + """ + pdb = pdb_dir / "1DAW.pdb" + if not pdb.exists(): + pytest.skip("1DAW.pdb fixture not present") + + from torchref.config import dtypes + from torchref.model import ModelFT + from torchref.refinement.targets.geometry.non_bonded import NonBondedTarget + + monkeypatch.setattr(dtypes, "_float", torch.float64) + model = ModelFT(device="cpu", dtype_float=torch.float64, verbose=0).load_pdb(str(pdb)) + target = NonBondedTarget(model) + + assert target.dtype_float == torch.float64 + for name in ("_sigma_vdw", "_r_exp", "_c_rep"): + assert getattr(target, name).dtype == torch.float64, name + + +@pytest.mark.unit +def test_register_target_warns_on_device_mismatch(pdb_dir): + """A mismatched target is reported, not silently relocated. + + Moving it would drag the model and data it borrows onto the state's + device -- registering a loss term must not have that side effect. + """ + pdb = pdb_dir / "1DAW.pdb" + if not pdb.exists(): + pytest.skip("1DAW.pdb fixture not present") + + from torchref.model import ModelFT + from torchref.refinement.loss_state import LossState + from torchref.refinement.targets.geometry.non_bonded import NonBondedTarget + + model = ModelFT(device="cpu", verbose=0).load_pdb(str(pdb)) + target = NonBondedTarget(model) + assert target.device.type == "cpu" + + state = LossState(device=torch.device("meta")) + with pytest.warns(UserWarning, match="is on cpu but the state is on"): + state.register_target("nb", target, probe=False) + + # ...and the target must not have been moved. + assert target.device.type == "cpu" + assert target._sigma_vdw.device.type == "cpu" diff --git a/tests/unit/test_device_mixin.py b/tests/unit/test_device_mixin.py index 5b52efc0..32009cab 100644 --- a/tests/unit/test_device_mixin.py +++ b/tests/unit/test_device_mixin.py @@ -248,3 +248,100 @@ def test_cell_cache_repopulates_on_target_dtype(): assert "volume" not in cell._cache, "cache must be cleared after .to()" new_volume = cell.volume assert new_volume.dtype == torch.float64 + + +# --------------------------------------------------------------------------- +# Probe-driven tracker inference (objects that own no tensors) +# --------------------------------------------------------------------------- + + +class _TensorFreeTracker(DeviceMixin, nn.Module): + """Carries device/dtype trackers but owns no tensor to read them from. + + This is the shape of every geometry target: the trackers say where the + object *will* allocate, and nothing it currently holds can confirm it. Such + objects are the only consumers of ``_probe_target``. + """ + + def __init__(self, dtype=torch.float64): + super().__init__() + self.device = torch.device("cpu") + self.dtype_float = dtype + + +@pytest.mark.unit +def test_device_move_does_not_clobber_dtype_tracker(any_device): + """A device-only move must leave ``dtype_float`` alone. + + Regression: the probe used a single float32 scratch for the dtype axis, so + a plain ``.to(device)`` reported ``float32`` and overwrote a float64 + tracker. Both axes need their own contrast pair. + """ + mod = _TensorFreeTracker(dtype=torch.float64) + mod.to(any_device) + assert mod.dtype_float == torch.float64, "device-only move changed dtype_float" + assert mod.device.type == any_device.type + + +@pytest.mark.unit +def test_dtype_only_move_does_not_clobber_device_tracker(): + """Mirror case: ``.float()`` must leave the device tracker alone.""" + mod = _TensorFreeTracker(dtype=torch.float64) + mod.device = torch.device("cpu") + mod.float() + assert mod.dtype_float == torch.float32 + assert mod.device == torch.device("cpu"), "dtype-only move changed device" + + +@pytest.mark.unit +@pytest.mark.parametrize( + "fn,expected_dtype", + [ + (lambda t: t.to("cpu"), None), + (lambda t: t.double(), torch.float64), + (lambda t: t.float(), torch.float32), + (lambda t: t.half(), torch.float16), + (lambda t: t, None), + ], + ids=["to_cpu", "double", "float", "half", "identity"], +) +def test_probe_target_dtype_axis(fn, expected_dtype): + """The dtype axis reports only what ``fn`` actually casts to. + + Host-independent: the dtype contrast pair is pinned to CPU, which is also + what keeps ``.double()`` probeable on Apple silicon (MPS has no float64, so + contrasting dtype on an accelerator scratch would fail outright). + """ + from torchref.utils.device_mixin import _probe_target + + _, dtype = _probe_target(fn) + assert dtype == expected_dtype + + +@pytest.mark.unit +@pytest.mark.gpu +@pytest.mark.parametrize( + "fn,preserves_device", + [ + (lambda t: t.to("cpu"), False), + (lambda t: t.double(), True), + (lambda t: t.float(), True), + (lambda t: t, True), + ], + ids=["to_cpu", "double", "float", "identity"], +) +def test_probe_target_device_axis(fn, preserves_device): + """The device axis distinguishes "moves to D" from "keeps its input's device". + + Requires an accelerator: telling those two apart needs two devices to + contrast. On a single-device host they are genuinely indistinguishable -- + and reporting ``cpu`` there is correct, since that is where everything is + anyway -- so there is nothing to assert. + """ + from torchref.utils.device_mixin import _probe_target + + device, _ = _probe_target(fn) + if preserves_device: + assert device is None, "device-preserving fn must not pin a device" + else: + assert device == torch.device("cpu") diff --git a/tests/unit/test_gradient_correctness.py b/tests/unit/test_gradient_correctness.py index e98eb081..17de1993 100644 --- a/tests/unit/test_gradient_correctness.py +++ b/tests/unit/test_gradient_correctness.py @@ -29,22 +29,14 @@ validated with the same cosine/gradnorm metric against central finite differences of a scalar loss (:func:`test_model_forward_xyz_adp_*`). -CPU-only by default; the Triton comparisons are ``gpu``/``cuda_only`` and skip -unless ``--run-gpu`` is passed on a CUDA host. +CPU-only by default; the Triton comparisons are marked ``cuda`` and are +auto-skipped unless the host actually has a CUDA device. """ -import itertools -import tempfile import pytest import torch -from torchref.base.direct_summation.dispatch import ( - _checkpointed_aniso, - _checkpointed_iso, - _eager_aniso, - _eager_iso, -) from torchref.base.targets.bond import _bond_math_eager, bond_math from torchref.base.targets.xray_gaussian import ( _gaussian_xray_loss_math_eager, @@ -54,123 +46,22 @@ pytestmark = pytest.mark.unit +# Kept for readability in the section header below; backend gating itself is +# the ``cuda`` marker's job (see tests/conftest.py). _HAS_CUDA = torch.cuda.is_available() # --------------------------------------------------------------------------- # Fixtures / helpers # --------------------------------------------------------------------------- -@pytest.fixture -def double_cpu(): - """Run the body in double precision on CPU, restoring global config after. - - ``gradcheck`` needs float64 leaves, and the eager SF reference casts ``hkl`` - to the configured ``dtypes.float`` internally, so the global dtype must be - float64 too. Pinning CPU keeps the eager paths off the Triton dispatch on - GPU hosts. - """ - f0, c0, d0 = dtypes.float, dtypes.complex, device.current - dtypes.float = torch.float64 - dtypes.complex = torch.complex128 - device.current = torch.device("cpu") - try: - yield - finally: - dtypes.float = f0 - dtypes.complex = c0 - device.current = d0 - - -def _flat_real(t: torch.Tensor) -> torch.Tensor: - """Flatten to a real 1-D vector, splitting complex into (re, im).""" - if t.is_complex(): - return torch.view_as_real(t).reshape(-1).double() - return t.reshape(-1).double() - - -def cosine_similarity(a: torch.Tensor, b: torch.Tensor) -> float: - """Cosine similarity of two gradients (direction agreement).""" - return torch.nn.functional.cosine_similarity( - _flat_real(a), _flat_real(b), dim=0 - ).item() - - -def gradnorm_ratio(a: torch.Tensor, b: torch.Tensor) -> float: - """Ratio of gradient norms ``||a|| / ||b||`` (magnitude agreement).""" - nb = _flat_real(b).norm() - return (_flat_real(a).norm() / nb).item() - - -def assert_grads_agree(got, ref, *, min_cos=0.9999, ratio_tol=1e-3, ctx=""): - """Assert per-leaf cosine ≈ 1 and gradnorm ratio ≈ 1 for matched grads.""" - for name, g, r in got_ref_pairs(got, ref): - cos = cosine_similarity(g, r) - ratio = gradnorm_ratio(g, r) - assert cos >= min_cos, f"{ctx}{name}: cosine {cos:.8f} < {min_cos}" - assert ( - abs(ratio - 1.0) <= ratio_tol - ), f"{ctx}{name}: gradnorm ratio {ratio:.8f} off 1 by > {ratio_tol}" - - -def got_ref_pairs(got, ref): - """Yield ``(name, got_grad, ref_grad)`` triples for a dict or sequence.""" - if isinstance(got, dict): - for k in got: - yield str(k), got[k], ref[k] - else: - for i, (g, r) in enumerate(zip(got, ref)): - yield str(i), g, r - - -# --- synthetic structure-factor inputs (P1, ~10-15 atoms) ------------------ -def _sf_inputs(N=12, R=18, dtype=torch.float64, device="cpu"): - g = torch.Generator().manual_seed(0) - hkl = torch.randint(-3, 4, (R, 3), generator=g).to(dtype=dtype, device=device) - s = (torch.rand(R, generator=g) * 0.4).to(dtype=dtype, device=device) - svec = (torch.randn(R, 3, generator=g) * 0.3).to(dtype=dtype, device=device) - A = torch.rand(N, 5, generator=g).to(dtype=dtype, device=device) - B = (torch.rand(N, 5, generator=g) + 0.5).to(dtype=dtype, device=device) - return hkl, s, svec, A, B - - -def _sf_leaves(N=12, dtype=torch.float64, device="cpu", requires_grad=True): - g = torch.Generator().manual_seed(1) - mk = lambda t: t.to(dtype=dtype, device=device).requires_grad_(requires_grad) - xyz = mk(torch.rand(N, 3, generator=g)) - occ = mk(torch.rand(N, generator=g) * 0.4 + 0.6) - adp = mk(torch.rand(N, generator=g) * 10 + 5) - U = mk(torch.rand(N, 6, generator=g) * 0.04 + 0.01) - return xyz, occ, adp, U - -# ============================================================================= -# 1. Structure-factor calculation — eager autograd vs finite differences -# ============================================================================= -def test_sf_iso_gradcheck(double_cpu): - """Isotropic SF: autograd matches numerical d/d(xyz, occ, adp).""" - hkl, s, _, A, B = _sf_inputs() - xyz, occ, adp, _ = _sf_leaves() - def f(x, o, a): - return _eager_iso(hkl, s, x, o, a, A, B, max_memory_gb=2.0) +# These live in ``tests/helpers/grad_asserts.py`` so the accelerator kernel +# comparison tests can share them; re-exported here because this module is where +# they originated and several tests below use them unqualified. +from tests.helpers.grad_asserts import assert_grads_agree # noqa: E402 - assert torch.autograd.gradcheck(f, (xyz, occ, adp), eps=1e-6, atol=1e-5) - - -def test_sf_aniso_gradcheck(double_cpu): - """Anisotropic SF: autograd matches numerical d/d(xyz, occ, U).""" - hkl, _, svec, A, B = _sf_inputs() - xyz, occ, _, U = _sf_leaves() - - def f(x, o, u): - return _eager_aniso(hkl, svec, x, o, u, A, B, max_memory_gb=2.0) - - assert torch.autograd.gradcheck(f, (xyz, occ, U), eps=1e-6, atol=1e-5) - - -# ============================================================================= -# 2. X-ray loss target — eager autograd vs finite differences -# ============================================================================= +# --- synthetic structure-factor inputs (P1, ~10-15 atoms) ------------------ def test_gaussian_xray_gradcheck(double_cpu): """Gaussian X-ray NLL (GaussianXrayTarget.forward body): d/d(F_obs, F_calc). @@ -225,196 +116,7 @@ def f(x): # 4. Optimized SF path (hand-written backward) vs eager autograd # Metric: cosine similarity + gradnorm ratio (CPU, always runs). # ============================================================================= -def _scalar(F): - """A non-trivial real scalar of a complex SF vector (mixes re & im).""" - return (F.real**2 + 2.0 * F.imag).sum() - - -def _grads(fn, leaves): - """Gradients of ``_scalar(fn(*leaves))`` w.r.t. each leaf.""" - out = _scalar(fn(*leaves)) - return torch.autograd.grad(out, leaves) - - -def test_checkpointed_iso_matches_eager_cosine(double_cpu): - """Isotropic ``_CheckpointedSF`` (recompute-on-backward) == eager autograd.""" - hkl, s, _, A, B = _sf_inputs() - xyz, occ, adp, _ = _sf_leaves() - xyz2, occ2, adp2, _ = _sf_leaves() - - g_ckpt = _grads( - lambda x, o, a: _checkpointed_iso(hkl, s, x, o, a, A, B, max_memory_gb=1e-7), - (xyz, occ, adp), - ) - g_eager = _grads( - lambda x, o, a: _eager_iso(hkl, s, x, o, a, A, B, max_memory_gb=2.0), - (xyz2, occ2, adp2), - ) - assert_grads_agree(g_ckpt, g_eager, ctx="iso ") - - -def test_checkpointed_aniso_matches_eager_cosine(double_cpu): - """Anisotropic ``_CheckpointedSF`` == eager autograd.""" - hkl, _, svec, A, B = _sf_inputs() - xyz, occ, _, U = _sf_leaves() - xyz2, occ2, _, U2 = _sf_leaves() - - g_ckpt = _grads( - lambda x, o, u: _checkpointed_aniso( - hkl, svec, x, o, u, A, B, max_memory_gb=1e-7 - ), - (xyz, occ, U), - ) - g_eager = _grads( - lambda x, o, u: _eager_aniso(hkl, svec, x, o, u, A, B, max_memory_gb=2.0), - (xyz2, occ2, U2), - ) - assert_grads_agree(g_ckpt, g_eager, ctx="aniso ") - - -# ============================================================================= -# 5. Production FFT path ModelFT.forward — autograd vs central finite diff -# Metric: cosine similarity + gradnorm ratio of a scalar loss. -# ============================================================================= -_P1_PDB_HEADER = ( - "CRYST1 20.000 20.000 20.000 90.00 90.00 90.00 P 1 1" -) - - -def _write_p1_pdb(n=12): - g = torch.Generator().manual_seed(0) - elems = "C N O S C N O C N O C S".split()[:n] - lines = [_P1_PDB_HEADER] - for i in range(n): - x, y, z = (torch.rand(3, generator=g) * 15 + 2.5).tolist() - e = elems[i] - lines.append( - f"ATOM {i + 1:5d} {e:<2s} GLY A{i + 1:4d} " - f"{x:8.3f}{y:8.3f}{z:8.3f} 1.00 20.00 {e:>2s}" - ) - lines.append("END") - f = tempfile.NamedTemporaryFile("w", suffix=".pdb", delete=False) - f.write("\n".join(lines) + "\n") - f.close() - return f.name - - -def _central_fd_grad(loss_fn, model, param, eps): - """Central finite-difference gradient of ``loss_fn()`` w.r.t. ``param``.""" - base = param.detach().clone() - flat = base.reshape(-1) - grad = torch.zeros_like(flat) - for i in range(flat.numel()): - for sign in (+1, -1): - pert = flat.clone() - pert[i] += sign * eps - with torch.no_grad(): - param.copy_(pert.view_as(base)) - model.reset_cache() - val = loss_fn() - if sign == +1: - vp = val - else: - vm = val - grad[i] = (vp - vm) / (2 * eps) - with torch.no_grad(): - param.copy_(base) - model.reset_cache() - return grad.view_as(base) - - -@pytest.mark.parametrize("which,eps", [("xyz", 5e-3), ("adp", 1e-2)]) -def test_model_forward_fft_gradient(which, eps): - """``ModelFT.forward`` (FFT path) autograd grad agrees with finite diff. - - Honors the paper's headline check — ``model.forward(hkl)`` w.r.t. ``xyz`` - and ``adp`` — for the production FFT engine. The FFT output is float32 and - the on-grid splat defeats element-wise gradcheck, so we compare a scalar - loss's autograd gradient against central finite differences via cosine - similarity and gradnorm ratio (robust where atol/rtol are not). - """ - from torchref.model.model_ft import ModelFT - - d0 = device.current - device.current = torch.device("cpu") # keep float32 default; avoid Triton - try: - model = ModelFT(max_res=1.5, verbose=0, device=torch.device("cpu")) - model.load_pdb(_write_p1_pdb()) - hkls = [h for h in itertools.product(range(-3, 4), repeat=3) if any(h)] - hkl = torch.tensor(hkls, dtype=torch.float32) - - def loss_fn(): - sf = model(hkl, recalc=True) - return (sf.real**2 + sf.imag**2).sum() - - param = getattr(model, which).refinable_params - g_auto = torch.autograd.grad(loss_fn(), param)[0] - g_fd = _central_fd_grad(loss_fn, model, param, eps) - - cos = cosine_similarity(g_auto, g_fd) - ratio = gradnorm_ratio(g_auto, g_fd) - assert cos >= 0.999, f"{which}: cosine {cos:.6f} < 0.999" - assert abs(ratio - 1.0) <= 0.02, f"{which}: gradnorm ratio {ratio:.6f}" - finally: - device.current = d0 - - -# ============================================================================= -# 6. Triton kernels (CUDA float32, hand-written backward) vs eager autograd -# Metric: cosine similarity + gradnorm ratio. Skipped without --run-gpu. -# ============================================================================= -@pytest.mark.gpu -@pytest.mark.cuda_only -@pytest.mark.skipif(not _HAS_CUDA, reason="Triton kernels require CUDA") -def test_triton_sf_iso_matches_eager_cosine(): - """DS isotropic Triton kernel gradient == eager autograd (CUDA float32).""" - from torchref.base.direct_summation.dispatch import ds_iso - from torchref.utils import Engine - - dev = "cuda" - hkl, s, _, A, B = _sf_inputs(dtype=torch.float32, device=dev) - xyz, occ, adp, _ = _sf_leaves(dtype=torch.float32, device=dev) - xyz2, occ2, adp2, _ = _sf_leaves(dtype=torch.float32, device=dev) - - g_triton = _grads( - lambda x, o, a: ds_iso(hkl, s, x, o, a, A, B, engine=Engine.TRITON), - (xyz, occ, adp), - ) - g_eager = _grads( - lambda x, o, a: _eager_iso(hkl, s, x, o, a, A, B, max_memory_gb=2.0), - (xyz2, occ2, adp2), - ) - # Looser thresholds: float32 + distinct kernel arithmetic. - assert_grads_agree(g_triton, g_eager, min_cos=0.999, ratio_tol=1e-2, ctx="iso ") - - -@pytest.mark.gpu -@pytest.mark.cuda_only -@pytest.mark.skipif(not _HAS_CUDA, reason="Triton kernels require CUDA") -def test_triton_sf_aniso_matches_eager_cosine(): - """DS anisotropic Triton kernel gradient == eager autograd (CUDA float32).""" - from torchref.base.direct_summation.dispatch import ds_aniso - from torchref.utils import Engine - - dev = "cuda" - hkl, _, svec, A, B = _sf_inputs(dtype=torch.float32, device=dev) - xyz, occ, _, U = _sf_leaves(dtype=torch.float32, device=dev) - xyz2, occ2, _, U2 = _sf_leaves(dtype=torch.float32, device=dev) - - g_triton = _grads( - lambda x, o, u: ds_aniso(hkl, svec, x, o, u, A, B, engine=Engine.TRITON), - (xyz, occ, U), - ) - g_eager = _grads( - lambda x, o, u: _eager_aniso(hkl, svec, x, o, u, A, B, max_memory_gb=2.0), - (xyz2, occ2, U2), - ) - assert_grads_agree(g_triton, g_eager, min_cos=0.999, ratio_tol=1e-2, ctx="aniso ") - - -@pytest.mark.gpu -@pytest.mark.cuda_only -@pytest.mark.skipif(not _HAS_CUDA, reason="Triton kernels require CUDA") +@pytest.mark.cuda def test_triton_gaussian_xray_matches_eager_cosine(): """Gaussian X-ray Triton kernel gradient == eager autograd (CUDA float32).""" dev = "cuda" @@ -435,9 +137,7 @@ def test_triton_gaussian_xray_matches_eager_cosine(): assert_grads_agree([g_t], [g_e], min_cos=0.999, ratio_tol=1e-2, ctx="xray ") -@pytest.mark.gpu -@pytest.mark.cuda_only -@pytest.mark.skipif(not _HAS_CUDA, reason="Triton kernels require CUDA") +@pytest.mark.cuda def test_triton_bond_matches_eager_cosine(): """Bond restraint Triton kernel gradient == eager autograd (CUDA float32).""" dev = "cuda" diff --git a/tests/unit/test_imports_smoke.py b/tests/unit/test_imports_smoke.py index d3800589..ed8f56bc 100644 --- a/tests/unit/test_imports_smoke.py +++ b/tests/unit/test_imports_smoke.py @@ -15,6 +15,7 @@ """ import importlib +import os import pkgutil import pytest @@ -95,3 +96,48 @@ def test_kinetic_subpackage_imports(): """ importlib.import_module("torchref.experimental.kinetic") importlib.import_module("torchref.experimental.kinetic.refinement") + + +@pytest.mark.unit +def test_import_survives_a_broken_triton_install(tmp_path): + """A Triton that *raises* on import must not take ``import torchref`` down with it. + + Absent Triton was always handled; a **broken** one was not. The optional-kernel guard + caught only ``ImportError``, so a Triton whose import raises anything else -- a driver + or LLVM version skew, which is the common real-world failure -- propagated straight out + of ``import torchref``, making the whole package unusable on a host whose GPU stack had + drifted. ``triton_available()`` was written to absorb exactly this and never got the + chance. + + Run in a subprocess with a deliberately-exploding ``triton`` shadowing the real one: + the failure happens at *import* time, so it cannot be provoked in-process once + ``torchref`` is already loaded. + """ + import subprocess + import sys + + pkg = tmp_path / "triton" + pkg.mkdir() + boom = 'raise RuntimeError("simulated driver/LLVM version skew")\n' + (pkg / "__init__.py").write_text(boom) + (pkg / "language.py").write_text(boom) + + env = dict(os.environ, PYTHONPATH=str(tmp_path)) + proc = subprocess.run( + [ + sys.executable, + "-c", + "import torchref\n" + "from torchref.utils import triton_available\n" + "assert triton_available() is False, 'a broken triton must read as absent'\n" + "print('ok')\n", + ], + capture_output=True, + text=True, + env=env, + ) + assert proc.returncode == 0, ( + "import torchref failed with a broken triton on the path:\n" + f"{proc.stderr[-2000:]}" + ) + assert "ok" in proc.stdout diff --git a/tests/unit/test_kernel_fixes.py b/tests/unit/test_kernel_fixes.py index ea59670d..362a53b7 100644 --- a/tests/unit/test_kernel_fixes.py +++ b/tests/unit/test_kernel_fixes.py @@ -23,51 +23,9 @@ pytestmark = pytest.mark.unit -@pytest.fixture -def double_cpu(): - """float64/complex128 on CPU for the duration of a test; restore after.""" - f0, c0, d0 = dtypes.float, dtypes.complex, device.current - dtypes.float = torch.float64 - dtypes.complex = torch.complex128 - device.current = torch.device("cpu") - try: - yield - finally: - dtypes.float = f0 - dtypes.complex = c0 - device.current = d0 - - # --------------------------------------------------------------------------- # Fix 5 — None scattering factors on the non-batched eager path # --------------------------------------------------------------------------- -def test_eager_sf_none_scattering_no_batch(double_cpu): - """`max_memory_gb=None` + `scattering_factors=None` computes from A/B.""" - from torchref.base.direct_summation.dispatch import _eager_aniso, _eager_iso - - g = torch.Generator().manual_seed(0) - N, R = 8, 12 - hkl = torch.randint(-3, 4, (R, 3), generator=g).double() - s = torch.rand(R, generator=g).double() * 0.4 - svec = torch.randn(R, 3, generator=g).double() * 0.3 - A = torch.rand(N, 5, generator=g).double() - B = torch.rand(N, 5, generator=g).double() + 0.5 - xyz = torch.rand(N, 3, generator=g).double() - occ = torch.rand(N, generator=g).double() * 0.4 + 0.6 - adp = torch.rand(N, generator=g).double() * 10 + 5 - U = torch.rand(N, 6, generator=g).double() * 0.04 + 0.01 - - f_none = _eager_iso(hkl, s, xyz, occ, adp, A, B, max_memory_gb=None) - f_batch = _eager_iso(hkl, s, xyz, occ, adp, A, B, max_memory_gb=2.0) - assert torch.isfinite(f_none).all() - assert torch.allclose(f_none, f_batch) - - a_none = _eager_aniso(hkl, svec, xyz, occ, U, A, B, max_memory_gb=None) - a_batch = _eager_aniso(hkl, svec, xyz, occ, U, A, B, max_memory_gb=2.0) - assert torch.isfinite(a_none).all() - assert torch.allclose(a_none, a_batch) - - # --------------------------------------------------------------------------- # Fix 3 — eager geometry math is NaN-safe at degenerate geometry # --------------------------------------------------------------------------- @@ -208,177 +166,3 @@ def test_forward_dtype_mismatch_raises(tmp_path): device.current = d0 -# --------------------------------------------------------------------------- -# CPU structured scatter: second derivatives must be CORRECT (not silently wrong) -# --------------------------------------------------------------------------- -def test_cpu_plain_scatter_second_derivative_correct(double_cpu, tmp_path): - """A Hessian-vector product through ``model.forward`` must match central - finite differences under a float64 config. - - Under ``double_cpu`` the density grid is float64, so the electron-density - dispatch routes to the portable plain ``scatter_add`` splat (the C++ fast - scatter is float32-only). This guards double-backward correctness of that - portable path -- the sanctioned float64 / EAGER Hessian route. The C++ fast - scatter's own double-backward is covered separately (float32) by - ``test_cpu_cpp_scatter_second_derivative_correct``. - """ - import itertools - - from torchref.model.model_ft import ModelFT - - pdb = tmp_path / "p1.pdb" - g = torch.Generator().manual_seed(0) - lines = ["CRYST1 20.000 20.000 20.000 90.00 90.00 90.00 P 1 1"] - for i in range(8): - x, y, z = (torch.rand(3, generator=g) * 15 + 2.5).tolist() - lines.append( - f"ATOM {i + 1:5d} C GLY A{i + 1:4d} " - f"{x:8.3f}{y:8.3f}{z:8.3f} 1.00 20.00 C" - ) - lines.append("END") - pdb.write_text("\n".join(lines) + "\n") - - model = ModelFT(max_res=2.0, verbose=0) - model.load_pdb(str(pdb)) - hkl = torch.tensor( - [h for h in itertools.product(range(-2, 3), repeat=3) if any(h)], - dtype=torch.float64, - ) - x = model.xyz.refinable_params - x0 = x.detach().clone() - gen = torch.Generator().manual_seed(1) - v = torch.randn(x0.shape, generator=gen) - v /= v.norm() - - def grad_at(xval): - with torch.no_grad(): - x.copy_(xval) - model.reset_cache() - sf = model(hkl, recalc=True) - return ( - torch.autograd.grad((sf.real**2 + sf.imag**2).sum(), x)[0].detach().clone() - ) - - # Autograd Hessian-vector product (double backward through the scatter). - with torch.no_grad(): - x.copy_(x0) - model.reset_cache() - sf = model(hkl, recalc=True) - (g1,) = torch.autograd.grad((sf.real**2 + sf.imag**2).sum(), x, create_graph=True) - (hvp_auto,) = torch.autograd.grad((g1 * v).sum(), x) - hvp_auto = hvp_auto.detach().clone() - - # Finite-difference reference. - eps = 3e-4 - hvp_fd = (grad_at(x0 + eps * v) - grad_at(x0 - eps * v)) / (2 * eps) - with torch.no_grad(): - x.copy_(x0) - model.reset_cache() - - cos = torch.nn.functional.cosine_similarity( - hvp_auto.flatten(), hvp_fd.flatten(), dim=0 - ).item() - rel = ((hvp_auto - hvp_fd).norm() / hvp_fd.norm()).item() - assert cos > 0.999, f"HVP direction wrong: cosine={cos:.6f}" - assert rel < 5e-3, f"HVP magnitude wrong: rel_err={rel:.3e}" - - -def test_cpu_cpp_scatter_second_derivative_correct(tmp_path): - """The CPU C++ fast-scatter double-backward must match central finite - differences (float32). - - Regression for the silent-wrong-Hessian bug: ``_StructuredScatterAdd`` - computed its backward with an imperative C++ gather (no ``grad_fn``), which - did not error under ``create_graph=True`` but dropped the second-order term - through the scatter transpose (measured cosine ~0.57 vs finite diff). The - fix routes the backward through the adjoint ``_StructuredGather`` Function - (scatter/gather are linear mutual adjoints), restoring a correct Hessian. - - This runs under a float32 config so the density dispatch stays on the C++ - fast scatter (the float64 config routes to the plain path instead -- see - ``test_cpu_plain_scatter_second_derivative_correct``). Finite differences - are float32 here, so the eps is larger and the tolerance looser than the - float64 plain-path test, but still far tighter than the ~0.57 regression. - """ - import itertools - - from torchref.base.electron_density.kernels.cpu.scatter_dispatch import ( - _get_cpp_scatter, - ) - from torchref.config import device, dtypes - from torchref.model.model_ft import ModelFT - - if _get_cpp_scatter() is None: - pytest.skip("C++ scatter extension not available; nothing to test here") - - f0, c0, d0 = dtypes.float, dtypes.complex, device.current - dtypes.float = torch.float32 - dtypes.complex = torch.complex64 - device.current = torch.device("cpu") - try: - pdb = tmp_path / "p1.pdb" - g = torch.Generator().manual_seed(0) - lines = ["CRYST1 20.000 20.000 20.000 90.00 90.00 90.00 P 1 1"] - for i in range(8): - x, y, z = (torch.rand(3, generator=g) * 15 + 2.5).tolist() - lines.append( - f"ATOM {i + 1:5d} C GLY A{i + 1:4d} " - f"{x:8.3f}{y:8.3f}{z:8.3f} 1.00 20.00 C" - ) - lines.append("END") - pdb.write_text("\n".join(lines) + "\n") - - model = ModelFT(max_res=2.0, verbose=0) - model.load_pdb(str(pdb)) - # The C++ fast scatter only runs on a float32 density grid. - assert model.xyz.refinable_params.dtype == torch.float32 - hkl = torch.tensor( - [h for h in itertools.product(range(-2, 3), repeat=3) if any(h)], - dtype=torch.float32, - ) - x = model.xyz.refinable_params - x0 = x.detach().clone() - gen = torch.Generator().manual_seed(1) - v = torch.randn(x0.shape, generator=gen) - v /= v.norm() - - def grad_at(xval): - with torch.no_grad(): - x.copy_(xval) - model.reset_cache() - sf = model(hkl, recalc=True) - return ( - torch.autograd.grad((sf.real**2 + sf.imag**2).sum(), x)[0] - .detach() - .clone() - ) - - # Autograd Hessian-vector product (double backward through the scatter). - with torch.no_grad(): - x.copy_(x0) - model.reset_cache() - sf = model(hkl, recalc=True) - (g1,) = torch.autograd.grad( - (sf.real**2 + sf.imag**2).sum(), x, create_graph=True - ) - (hvp_auto,) = torch.autograd.grad((g1 * v).sum(), x) - hvp_auto = hvp_auto.detach().clone() - - # Central finite-difference reference. Larger eps than the float64 test: - # in float32 a too-small step is swamped by subtraction noise. - eps = 3e-3 - hvp_fd = (grad_at(x0 + eps * v) - grad_at(x0 - eps * v)) / (2 * eps) - with torch.no_grad(): - x.copy_(x0) - model.reset_cache() - - cos = torch.nn.functional.cosine_similarity( - hvp_auto.flatten(), hvp_fd.flatten(), dim=0 - ).item() - rel = ((hvp_auto - hvp_fd).norm() / hvp_fd.norm()).item() - assert cos > 0.9999, f"HVP direction wrong: cosine={cos:.6f}" - assert rel < 5e-3, f"HVP magnitude wrong: rel_err={rel:.3e}" - finally: - dtypes.float = f0 - dtypes.complex = c0 - device.current = d0 diff --git a/tests/unit/test_seeded_lbfgs.py b/tests/unit/test_seeded_lbfgs.py new file mode 100644 index 00000000..83147f48 --- /dev/null +++ b/tests/unit/test_seeded_lbfgs.py @@ -0,0 +1,199 @@ +"""Unit tests for SeededLBFGS and the Hessian-diagonal curvature helper. + +Covers: +1. A diagonal quadratic bowl: an exact inverse-diagonal seed makes the first + (single) inner iteration land on the minimum, where plain steepest-descent + L-BFGS with the same one-iteration budget does not. +2. ``init_hess_diag=None`` reproduces ``torch.optim.LBFGS`` bit-for-bit; an + all-ones seed reproduces the first *direction* but not the first step length. +3. The Hutchinson diagonal helper matches the true Hessian diagonal (exact + ``probe="basis"``; statistical ``probe="rademacher"``). + +All tests are pure-CPU (the analytic double-backward HVP runs under +``use_portable()`` on plain torch), so no GPU marker is needed. +""" + +import pytest +import torch + +from torchref.refinement.optimizers import ( + SeededLBFGS, + hessian_diagonal, + preconditioner_from_diagonal, +) + +pytestmark = pytest.mark.unit + + +# ============================================================================= +# 1. Diagonal quadratic bowl — one seeded step reaches the minimum +# ============================================================================= +def test_seeded_step_solves_diagonal_quadratic_in_one_iter(): + """f(x) = 0.5 * sum(h_i x_i^2); seed = 1/h => d0 = -x, t = 1 => x -> 0.""" + h = torch.tensor([1.0, 4.0, 9.0, 16.0], dtype=torch.float64) + x0 = torch.tensor([3.0, -2.0, 1.0, 0.5], dtype=torch.float64) + + # --- seeded --- + x = torch.nn.Parameter(x0.clone()) + opt = SeededLBFGS( + [x], lr=1.0, max_iter=1, history_size=100, line_search_fn="strong_wolfe" + ) + opt.set_init_hess_diag(1.0 / h) # exact inverse-diagonal (lambda = 0) + + def closure(): + opt.zero_grad() + loss = 0.5 * (h * x * x).sum() + loss.backward() + return loss + + opt.step(closure) + assert x.detach().abs().max().item() < 1e-8 + + # --- stock steepest descent, same one-iteration budget: does NOT solve --- + x2 = torch.nn.Parameter(x0.clone()) + stock = torch.optim.LBFGS( + [x2], lr=1.0, max_iter=1, history_size=100, line_search_fn="strong_wolfe" + ) + + def closure2(): + stock.zero_grad() + loss = 0.5 * (h * x2 * x2).sum() + loss.backward() + return loss + + stock.step(closure2) + # one steepest-descent line search along -g cannot zero all four coords + assert x2.detach().abs().max().item() > 1e-2 + + +# ============================================================================= +# 2. Equivalence / divergence vs stock torch.optim.LBFGS +# ============================================================================= +def _quartic_closure(opt, x, a): + def closure(): + opt.zero_grad() + loss = ((x - a) ** 4).sum() + loss.backward() + return loss + + return closure + + +def test_none_seed_is_bit_identical_to_stock_lbfgs(): + """With init_hess_diag=None the whole trajectory matches torch.optim.LBFGS.""" + a = torch.tensor([1.0, -2.0, 0.5, 3.0], dtype=torch.float64) + x0 = torch.zeros(4, dtype=torch.float64) + + kw = dict(lr=1.0, max_iter=20, history_size=10, line_search_fn="strong_wolfe") + + x_seed = torch.nn.Parameter(x0.clone()) + opt_seed = SeededLBFGS([x_seed], **kw) # no seed set -> None + c_seed = _quartic_closure(opt_seed, x_seed, a) + + x_stock = torch.nn.Parameter(x0.clone()) + opt_stock = torch.optim.LBFGS([x_stock], **kw) + c_stock = _quartic_closure(opt_stock, x_stock, a) + + for _ in range(5): + l_seed = float(opt_seed.step(c_seed)) + l_stock = float(opt_stock.step(c_stock)) + assert l_seed == l_stock + assert torch.equal(x_seed.detach(), x_stock.detach()) + + +def test_ones_seed_matches_stock_direction_but_not_step_length(): + """An all-ones seed gives the stock first direction (-g) but t = lr, not the + ``min(1, 1/||g||_1)`` heuristic.""" + a = torch.tensor([1.0, -2.0, 0.5, 3.0], dtype=torch.float64) + x0 = torch.zeros(4, dtype=torch.float64) + kw = dict(lr=1.0, max_iter=1, history_size=10, line_search_fn="strong_wolfe") + + x_seed = torch.nn.Parameter(x0.clone()) + opt_seed = SeededLBFGS([x_seed], **kw) + opt_seed.set_init_hess_diag(torch.ones(4, dtype=torch.float64)) + opt_seed.step(_quartic_closure(opt_seed, x_seed, a)) + + x_stock = torch.nn.Parameter(x0.clone()) + opt_stock = torch.optim.LBFGS([x_stock], **kw) + opt_stock.step(_quartic_closure(opt_stock, x_stock, a)) + + d_seed = opt_seed.state[opt_seed._params[0]]["d"] + d_stock = opt_stock.state[opt_stock._params[0]]["d"] + # same first direction (steepest descent, since precond = ones) + assert torch.allclose(d_seed, d_stock) + # but the initial trial step length differs by design + t_seed = opt_seed.state[opt_seed._params[0]]["t"] + t_stock = opt_stock.state[opt_stock._params[0]]["t"] + # (post-line-search t may coincide only by accident; assert the *first-guess* + # rule differs by checking the resulting parameter points are not identical) + assert not torch.equal(x_seed.detach(), x_stock.detach()) or t_seed != t_stock + + +# ============================================================================= +# 3. Hessian-diagonal helper vs the true diagonal +# ============================================================================= +def test_hessian_diagonal_basis_matches_true_diagonal(): + """Exact (basis-probe) diagonal matches torch.autograd.functional.hessian.""" + g = torch.Generator().manual_seed(0) + n = 6 + M = torch.randn(n, n, generator=g, dtype=torch.float64) + M = 0.5 * (M + M.t()) # symmetric (off-diagonal curvature present) + c = torch.rand(n, generator=g, dtype=torch.float64) + 0.1 + x_val = torch.randn(n, generator=g, dtype=torch.float64) + + def f(v): + return 0.5 * (v @ M @ v) + (c * v**4).sum() + + x = torch.nn.Parameter(x_val.clone()) + + diag_est = hessian_diagonal(lambda: f(x), [x], probe="basis") + + H = torch.autograd.functional.hessian(f, x_val.clone()) + true_diag = torch.diagonal(H) + assert torch.allclose(diag_est, true_diag, atol=1e-8) + + +def test_per_group_floor_does_not_crush_small_group(): + """A two-group diagonal with a 1000:1 magnitude gap: the global floor crushes + the small group's preconditioner, but the per-group floor gives each group its + own Newton scale.""" + # group A: large curvature (~1e3), group B: small curvature (~1.0) + diag = torch.tensor([1000.0, 800.0, 1.0, 0.8], dtype=torch.float64) + sizes = [2, 2] + lam = 1e-2 + + glob = preconditioner_from_diagonal(diag, lam=lam) # global floor + per = preconditioner_from_diagonal(diag, lam=lam, group_sizes=sizes) + + # Global floor = lam*max|diag| = 10 -> small group clamped to 1/10 = 0.1, + # i.e. its true 1/1.0=1.0 preconditioner is crushed ~10x. + assert torch.allclose(glob[2:], torch.tensor([0.1, 0.1], dtype=torch.float64)) + # Per-group: small group floored by its OWN max (1.0) -> ~1/|diag| preserved. + assert per[2] > 0.9 # ~1/1.0, not crushed + assert per[3] > 0.9 # 1/0.8 clamped by its own floor 1e-2*1.0 -> ~1/0.8=1.25 + # large group identical either way (it sets the global max) + assert torch.allclose(per[:2], glob[:2]) + + +def test_hessian_diagonal_rademacher_approximates_true_diagonal(): + """Stochastic (Rademacher) diagonal converges to the truth within tolerance.""" + g = torch.Generator().manual_seed(1) + n = 6 + M = torch.randn(n, n, generator=g, dtype=torch.float64) + M = 0.5 * (M + M.t()) + c = torch.rand(n, generator=g, dtype=torch.float64) + 0.1 + x_val = torch.randn(n, generator=g, dtype=torch.float64) + + def f(v): + return 0.5 * (v @ M @ v) + (c * v**4).sum() + + x = torch.nn.Parameter(x_val.clone()) + + probe_gen = torch.Generator().manual_seed(2) + diag_hutch = hessian_diagonal( + lambda: f(x), [x], n_probes=4000, generator=probe_gen, probe="rademacher" + ) + + true_diag = torch.diagonal(torch.autograd.functional.hessian(f, x_val.clone())) + rel = (diag_hutch - true_diag).norm() / true_diag.norm() + assert rel < 0.1 diff --git a/tests/unit/utils/test_backend_tables.py b/tests/unit/utils/test_backend_tables.py new file mode 100644 index 00000000..1e1cbfa2 --- /dev/null +++ b/tests/unit/utils/test_backend_tables.py @@ -0,0 +1,184 @@ +"""The three production dispatch tables, checked as policy rather than as numerics. + +Separate from ``test_backends.py``, which exercises the resolver against synthetic tables. +This file asserts things about the tables that actually ship: that each backend is available +on the hosts where it is declared to be, that each row resolves to the kernel its name claims, +and that selection returns what the docs say for every reachable ``(device, dtype)``. + +Everything here runs on any host. That is the point -- the accelerator provenance tests can +only cover the backend the machine happens to have, whereas ``select`` is a pure function of +the table and can be interrogated for combinations this host cannot execute. +""" + +import pytest +import torch + +from torchref.base.direct_summation._backends import DS_BACKENDS +from torchref.base.electron_density._backends import DENSITY_BACKENDS +from torchref.base.targets._dispatch import TARGET_BACKENDS +from torchref.utils.backends import select + +pytestmark = pytest.mark.unit + +ALL_TABLES = (DENSITY_BACKENDS, DS_BACKENDS, TARGET_BACKENDS) + +ALL_BACKENDS = [ + pytest.param(t, b, id=f"{t.name.split()[0]}-{b.name}") + for t in ALL_TABLES + for b in t.backends +] + + + + +# --------------------------------------------------------------------------- +# Availability: declared expectations must match reality +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("table,backend", ALL_BACKENDS) +def test_backend_is_available_where_it_is_expected(table, backend): + """A backend declared required on this host must actually be usable here. + + One parametrized test in place of a bespoke compile check per kernel -- that is what + ``expect_available`` on the table row is for. + + It **fails rather than skips**, and that asymmetry is deliberate. Dispatch degrades + quietly, which is correct for users and dangerous for CI: if every test skipped when a + kernel went missing, a broken build would produce an all-green run while production had + silently fallen back to a ~100x slower path. + + This is also what keeps the accelerator-vs-reference comparisons honest now that forcing + an engine is no longer possible. A skewed Triton install used to be caught by + ``Engine.TRITON`` raising; it is now caught here, on any run on that host, which is + strictly earlier and louder. + """ + if not backend.expected_here(): + pytest.skip(f"{backend.name} is not expected on this host") + why = backend.unavailable() + assert why is None, ( + f"{table.name}/{backend.name} is expected to work on this host but is not " + f"available: {why}" + ) + + +# --------------------------------------------------------------------------- +# Provenance without a call recorder +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("table", ALL_TABLES, ids=lambda t: t.name.split()[0]) +def test_every_row_resolves_to_a_distinct_kernel(table): + """No two rows may resolve to the same function, and iso must differ from aniso. + + Catches the one failure mode nothing else can: a row naming the *wrong but valid* kernel. + If ``cpu_sphere`` pointed at ``add_isotropic_plain_var``, the accuracy tests would pass + (portable *is* the reference implementation), the selection tests would pass (the name it + returns is still right), and the availability test would pass (the extension built fine) -- + while the CPU fast path was silently gone. Here it collides with ``portable`` and fails. + + Asserting distinctness rather than comparing against the declared attribute is what makes + this non-tautological: re-reading ``backend.kernel`` would only prove ``getattr`` works. + + Runs on every host, unlike the monkeypatched provenance tests, because resolving a kernel + imports its module without launching it -- the CUDA module is guarded on ``_HAVE_TRITON`` + and the Metal one defers shader compilation to first use. + """ + seen = {} + for backend in table.backends: + if backend.kernel is None: # gate-only row; the call site owns the kernel + continue + iso, aniso = backend.resolve(False), backend.resolve(True) + assert iso is not aniso, ( + f"{table.name}/{backend.name} resolves iso and aniso to the same function " + f"({iso.__module__}.{iso.__name__})" + ) + for variant, fn in (("iso", iso), ("aniso", aniso)): + key = (fn.__module__, fn.__name__) + if key in seen: + other, other_variant = seen[key] + pytest.fail( + f"{table.name}: {backend.name}/{variant} and {other}/{other_variant} " + f"both resolve to {key[0]}.{key[1]} -- one of these rows names the " + "wrong kernel, and every other test would still pass" + ) + seen[key] = (backend.name, variant) + + +# --------------------------------------------------------------------------- +# Selection, enumerated +# --------------------------------------------------------------------------- +def _six(device, dtype): + return [torch.zeros(4, device=device, dtype=dtype)] * 6 + + +_CPU_CASES = [ + (torch.float32, False, "cpu_sphere"), + (torch.float64, False, "cpu_sphere"), + (torch.float32, True, "portable"), + (torch.float64, True, "portable"), +] + + +@pytest.mark.parametrize( + "dtype,pin,expected", + _CPU_CASES, + ids=[f"{d}-{'portable' if p else 'default'}" for d, p, _ in _CPU_CASES], +) +def test_density_selection_on_cpu(dtype, pin, expected): + """The density policy, asserted directly instead of inferred from a call recorder. + + ``select`` is a pure function of the table, so this replaces most of what previously + needed a monkeypatched spy -- and unlike a spy it can assert the float64 legs, which the + accelerator provenance tests cannot reach at all. + """ + got = select(DENSITY_BACKENDS, _six("cpu", dtype), force_portable=pin).name + assert got == expected + + +@pytest.mark.mps +@pytest.mark.parametrize("pin,expected", [(False, "mps_metal"), (True, "portable")], + ids=["default", "portable"]) +def test_density_selection_on_mps(pin, expected): + got = select(DENSITY_BACKENDS, _six("mps", torch.float32), force_portable=pin).name + assert got == expected + + +@pytest.mark.parametrize("table", ALL_TABLES, ids=lambda t: t.name.split()[0]) +def test_force_portable_reaches_the_base_case_in_every_table(table): + """The override is uniform: one flag, same meaning at every dispatch site. + + Not true of what it replaced -- ``Engine.METAL`` selected the Metal splat at the density + site and silently meant "run eager" at the other two, a per-site asymmetry that existed + only as an early ``return False`` inside a predicate. + """ + ts = _six("cpu", torch.float32) + assert select(table, ts, force_portable=True) is table.base + + +def test_ds_selection_ignores_integer_hkl(): + """An ``int32`` ``hkl`` must not change what direct summation selects. + + Miller indices are ``int32`` in production and every DS kernel casts them itself, + exactly. ``hkl`` is outside the probe set for that reason, and the dtype rule exempts + integers regardless -- two independent guards against the production dtype reading as a + capability failure. + """ + n = 4 + f32 = lambda: torch.zeros(n, dtype=torch.float32) # noqa: E731 + args_int = [torch.zeros(n, 3, dtype=torch.int32), f32(), torch.zeros(n, 3)] + [f32()] * 4 + args_f32 = [torch.zeros(n, 3, dtype=torch.float32), f32(), torch.zeros(n, 3)] + [f32()] * 4 + assert ( + select(DS_BACKENDS, args_int).name == select(DS_BACKENDS, args_f32).name + ) + + +def test_second_order_capability_is_declared_not_guessed(): + """``second_order`` is the single source for which kernels compose under ``create_graph``. + + ``test_second_order.py`` derives its skip list from this, so a kernel gaining or losing a + double-backward path cannot leave the test matrix stale. + """ + density = {b.name: b.second_order for b in DENSITY_BACKENDS.backends} + assert density == { + "cuda_triton": False, + "mps_metal": False, + "cpu_sphere": True, + "portable": True, + } diff --git a/tests/unit/utils/test_backends.py b/tests/unit/utils/test_backends.py new file mode 100644 index 00000000..833e8ecc --- /dev/null +++ b/tests/unit/utils/test_backends.py @@ -0,0 +1,440 @@ +"""The dispatch resolver, tested against synthetic tables rather than real kernels. + +Synthetic on purpose: this file is about the *policy* machinery -- device/dtype matching, +probe ordering, failure handling, the portable override -- and a synthetic table can express +cases no real host can reach (a CUDA row on a CPU-only machine, a probe that raises if called +at all). The shipping tables are checked where they live, against the kernels they name, in +``test_backend_tables.py``. + +The three ``use_portable`` tests at the top are what survived the deletion of +``test_triton_dispatch.py``. Of its nineteen tests, only its state-management ones had a +subject that outlived the ``Engine`` enum; everything else asserted behaviour of forcing an +*accelerator*, which is no longer a concept, or device/dtype rules now asserted directly +against the shipping tables. +""" + +import pytest +import torch + +from torchref.utils.backends import ( + Backend, + BackendTable, + TorchRefDegradationWarning, + force_portable, + run_or_degrade, + select, + set_force_portable, + use_portable, + will_use, +) + +pytestmark = pytest.mark.unit + +_THIS = "tests.unit.utils.test_backends" + + +# --- kernels the synthetic tables point at ------------------------------- +def accel_iso(*a, **k): + return "accel_iso" + + +def accel_aniso(*a, **k): + return "accel_aniso" + + +def base_iso(*a, **k): + return "base_iso" + + +def base_aniso(*a, **k): + return "base_aniso" + + +def boom(*a, **k): + raise RuntimeError("kernel exploded") + + +def probe_ok(): + return None + + +def probe_missing(): + return "the widget is not installed" + + +def probe_must_not_run(): + raise AssertionError("availability probed before device/dtype was checked") + + +def _table( + *, + accel_probe=(_THIS, "probe_ok"), + accel_on_failure="degrade", + accel_device="cuda", +): + return BackendTable( + name="synthetic", + backends=( + Backend( + name="accel", + kernel=(_THIS, "accel_iso", "accel_aniso"), + device=accel_device, + dtypes=(torch.float32,), + probe=accel_probe, + on_failure=accel_on_failure, + second_order=False, + ), + Backend(name="base", kernel=(_THIS, "base_iso", "base_aniso")), + ), + ) + + +# --------------------------------------------------------------------------- +# The portable override +# --------------------------------------------------------------------------- +def test_force_portable_roundtrip(): + previous = force_portable() + try: + set_force_portable(True) + assert force_portable() is True + set_force_portable(False) + assert force_portable() is False + finally: + set_force_portable(previous) + + +def test_use_portable_nests_and_restores(): + previous = force_portable() + try: + set_force_portable(False) + with use_portable(): + assert force_portable() is True + with use_portable(): + assert force_portable() is True + assert force_portable() is True + assert force_portable() is False + finally: + set_force_portable(previous) + + +def test_use_portable_restores_on_exception(): + """A raise inside the block must not leave dispatch pinned for the rest of the process. + + Worth its own test because the failure is invisible and contagious: every later call in + the process would silently run the reference kernel, and a benchmark would report a + regression with no obvious cause. + """ + previous = force_portable() + try: + set_force_portable(False) + with pytest.raises(ValueError): + with use_portable(): + raise ValueError("boom") + assert force_portable() is False + finally: + set_force_portable(previous) + + +def test_force_portable_selects_the_base_case_over_a_matching_accelerator(): + """The override's whole job: pin the reference even when something faster would run. + + This is the one thing automatic selection cannot do for you. A fallback covers an + accelerator that is missing or throws; it cannot cover one that runs and returns *wrong + numbers*, which is what you are checking for when you reach for this. + """ + t = _table(accel_device="cpu") + ts = [torch.zeros(2)] + assert select(t, ts).name == "accel" + assert select(t, ts, force_portable=True).name == "base" + with use_portable(): + assert select(t, ts).name == "base" + + +def test_per_call_argument_overrides_the_process_wide_setting(): + previous = force_portable() + try: + t = _table(accel_device="cpu") + ts = [torch.zeros(2)] + set_force_portable(True) + assert select(t, ts).name == "base" + assert select(t, ts, force_portable=False).name == "accel" + finally: + set_force_portable(previous) + + +# --------------------------------------------------------------------------- +# Table invariants, asserted at construction +# --------------------------------------------------------------------------- +def test_table_requires_exactly_one_base_case(): + """Totality depends on it: with no unrestricted row, some inputs would match nothing. + + ``select`` has no error path, and that is structural rather than optimistic -- so the + condition that makes it true is checked where the table is defined. + """ + with pytest.raises(ValueError, match="exactly one unrestricted base"): + BackendTable( + name="two-bases", + backends=( + Backend("a", (_THIS, "base_iso", "base_aniso")), + Backend("b", (_THIS, "base_iso", "base_aniso")), + ), + ) + + +def test_backend_rejects_an_unknown_failure_policy(): + with pytest.raises(ValueError, match="on_failure"): + Backend("x", (_THIS, "base_iso", "base_aniso"), on_failure="explode") + + +def test_backend_rejects_an_unknown_expectation(): + with pytest.raises(ValueError, match="expect_available"): + Backend("x", (_THIS, "base_iso", "base_aniso"), expect_available="on tuesdays") + + +def test_expectation_conditions_are_evaluated_against_the_host(): + """``expect_available`` is a host predicate, not a static flag.""" + mk = lambda exp: Backend( # noqa: E731 + "b", (_THIS, "base_iso", "base_aniso"), expect_available=exp + ) + assert mk("always").expected_here() is True + assert mk("never").expected_here() is False + assert mk("cuda").expected_here() is torch.cuda.is_available() + + +def test_expectation_does_not_influence_selection(): + """A backend that "should" work still loses to an honest probe saying it does not. + + Keeping these separate is the point: ``expect_available`` is what CI asserts, ``probe`` + is what dispatch obeys. Conflating them would route real work to a kernel that cannot + run, on the strength of a mis-declared expectation. + """ + t = BackendTable( + name="expectation-vs-fact", + backends=( + Backend( + "accel", + (_THIS, "accel_iso", "accel_aniso"), + device="cpu", + probe=(_THIS, "probe_missing"), + expect_available="always", + ), + Backend("base", (_THIS, "base_iso", "base_aniso")), + ), + ) + assert t.by_name("accel").expected_here() is True + assert select(t, [torch.zeros(2)]).name == "base" + + +# --------------------------------------------------------------------------- +# Selection +# --------------------------------------------------------------------------- +def test_falls_through_to_the_base_case_on_a_wrong_device(): + assert select(_table(), [torch.zeros(2)]).name == "base" + + +def test_selects_the_accelerator_when_everything_matches(): + """Device match is expressed through the table, not the host, so this runs anywhere.""" + assert select(_table(accel_device="cpu"), [torch.zeros(2)]).name == "accel" + + +def test_selection_never_fails(): + """There is no input for which nothing is selected. + + The base case declares no device and no dtype and carries no probe, so it matches + unconditionally. Asserted against deliberately hostile inputs -- wrong device, wrong + dtype, an unavailable accelerator, and no tensors at all. + """ + t = _table(accel_probe=(_THIS, "probe_missing")) + for ts in ( + [torch.zeros(2)], + [torch.zeros(2, dtype=torch.float64)], + [None, None], + [], + ): + assert select(t, ts).name == "base" + + +def test_unavailability_falls_through(): + t = _table(accel_probe=(_THIS, "probe_missing"), accel_device="cpu") + assert select(t, [torch.zeros(2)]).name == "base" + + +def test_none_entries_are_ignored(): + t = _table(accel_device="cpu") + assert select(t, [None, torch.zeros(2), None]).name == "accel" + + +def test_probes_restrict_which_arguments_carry_the_contract(): + """A tensor outside ``probes`` must not affect selection. + + This is what lets direct summation police the refinable leaves while ignoring ``hkl``, + whose dtype provably costs nothing. + """ + table = BackendTable( + name="probed", + backends=( + Backend( + "accel", + (_THIS, "accel_iso", "accel_aniso"), + device="cpu", + dtypes=(torch.float32,), + probes=(0,), + ), + Backend("base", (_THIS, "base_iso", "base_aniso")), + ), + ) + f32, f64 = torch.zeros(2), torch.zeros(2, dtype=torch.float64) + assert select(table, [f32, f64]).name == "accel" # position 1 unprobed + assert select(table, [f64, f32]).name == "base" # position 0 still bites + + +def test_will_use_names_the_selected_backend(): + t = _table(accel_device="cpu") + ts = [torch.zeros(2)] + assert will_use(t, "accel", ts) is True + assert will_use(t, "base", ts) is False + assert will_use(t, "base", ts, force_portable=True) is True + + +# --------------------------------------------------------------------------- +# Two-phase ordering +# --------------------------------------------------------------------------- +def test_availability_is_not_probed_when_device_already_mismatches(): + """Device/dtype must be rejected before the probe is consulted. + + Enforced with a probe that raises if reached, because getting the order wrong is + invisible to a return-value assertion: an MPS host would compile the CPU C++ extension it + never uses, and a CPU-only host would import Triton to answer a question about a CPU + tensor. + """ + t = _table(accel_probe=(_THIS, "probe_must_not_run")) + assert select(t, [torch.zeros(2)]).name == "base" + + +def test_a_broken_probe_import_is_a_reason_not_a_crash(): + t = _table(accel_probe=("torchref.no.such.module", "nope"), accel_device="cpu") + assert select(t, [torch.zeros(2)]).name == "base" + assert "could not be imported" in t.by_name("accel").unavailable() + + +# --------------------------------------------------------------------------- +# dtype rules +# --------------------------------------------------------------------------- +def test_require_uniform_dtype_rejects_a_mixed_set_that_membership_would_admit(): + """The distinction is memory safety, not taste. + + ``dtypes=(f32, f64)`` read as membership admits a float64 map beside float32 atoms, which + the fused CPU kernel would reinterpret through a raw pointer of the map's type. + """ + mixed = BackendTable( + name="uniform", + backends=( + Backend( + "fused", + (_THIS, "accel_iso", "accel_aniso"), + device="cpu", + dtypes=(torch.float32, torch.float64), + require_uniform_dtype=True, + ), + Backend("base", (_THIS, "base_iso", "base_aniso")), + ), + ) + f32, f64 = torch.zeros(2), torch.zeros(2, dtype=torch.float64) + assert select(mixed, [f32, f32]).name == "fused" + assert select(mixed, [f64, f64]).name == "fused" + assert select(mixed, [f64, f32]).name == "base" + + +def test_integer_tensors_are_exempt_from_the_dtype_contract(): + """Integer inputs must not disqualify a float32 backend. + + Miller indices are ``int32`` in production and every kernel casts them itself, exactly + and losslessly. An identity test against float32 would reject the production dtype and + disable the very kernel the rule protects, so the contract applies to floats only. + """ + t = _table(accel_device="cpu") + hkl_int = torch.zeros(4, 3, dtype=torch.int32) + assert select(t, [hkl_int, torch.zeros(2)]).name == "accel" + # a float in the wrong precision still disqualifies + assert select(t, [hkl_int, torch.zeros(2, dtype=torch.float64)]).name == "base" + + +# --------------------------------------------------------------------------- +# Running +# --------------------------------------------------------------------------- +def test_run_resolves_the_kernel_late_so_it_stays_patchable(monkeypatch): + """The kernel is looked up at call time, at its defining module. + + Load-bearing: the provenance tests prove a named kernel really ran by patching it. A + table that captured function objects at import would leave those tests green while + measuring nothing. + """ + import tests.unit.utils.test_backends as self_mod + + t = _table(accel_device="cpu") + monkeypatch.setattr(self_mod, "accel_iso", lambda *a, **k: "patched") + assert run_or_degrade(t, t.by_name("accel"), False) == "patched" + + +def test_run_picks_the_anisotropic_variant(): + t = _table(accel_device="cpu") + b = t.by_name("accel") + assert run_or_degrade(t, b, False) == "accel_iso" + assert run_or_degrade(t, b, True) == "accel_aniso" + + +def test_gate_only_backend_cannot_be_run(): + """A row with ``kernel=None`` declares criteria for a call site that names its own kernel. + + The geometry targets are the case. Selecting such a row is fine; running it is a + programming error and says so, rather than failing with an opaque unpack. + """ + t = BackendTable( + name="gate-only", + backends=( + Backend("gate", kernel=None, device="cpu"), + Backend("base", kernel=(_THIS, "base_iso", "base_aniso")), + ), + ) + with pytest.raises(TypeError, match="gate-only"): + run_or_degrade(t, t.by_name("gate"), False) + + +def test_degrade_backend_falls_back_and_warns(monkeypatch): + """The fallback runs -- and says so. + + Asserted with ``pytest.warns`` because a silent fallback is the failure mode the warning + exists to prevent: both pytest configs promote ``TorchRefDegradationWarning`` to an error, + so a degradation anywhere else in the suite fails the run. This is the one place it is + expected, so it opts in. + """ + import tests.unit.utils.test_backends as self_mod + + t = _table(accel_device="cpu", accel_on_failure="degrade") + monkeypatch.setattr(self_mod, "accel_iso", boom) + with pytest.warns(TorchRefDegradationWarning, match="kernel exploded"): + got = run_or_degrade(t, t.by_name("accel"), False) + assert got == "base_iso" + + +def test_raise_backend_does_not_fall_back(monkeypatch): + """A production kernel that built fine and then threw is a bug, not a capability miss. + + Degrading would convert a wrong-results bug into a large silent slowdown whose output + still looks plausible, because the fallback implements the same contract. + """ + import tests.unit.utils.test_backends as self_mod + + t = _table(accel_device="cpu", accel_on_failure="raise") + monkeypatch.setattr(self_mod, "accel_iso", boom) + with pytest.raises(RuntimeError, match="kernel exploded"): + run_or_degrade(t, t.by_name("accel"), False) + + +def test_base_case_failure_always_propagates(monkeypatch): + import tests.unit.utils.test_backends as self_mod + + t = _table() + monkeypatch.setattr(self_mod, "base_iso", boom) + with pytest.raises(RuntimeError, match="kernel exploded"): + run_or_degrade(t, t.base, False) diff --git a/tests/unit/utils/test_device_resolution.py b/tests/unit/utils/test_device_resolution.py index 5ef428aa..7e5baa28 100644 --- a/tests/unit/utils/test_device_resolution.py +++ b/tests/unit/utils/test_device_resolution.py @@ -96,7 +96,7 @@ def test_explicit_overrides_skips_nones(self): # GPU tests (cuda required for the cross-device branches) # --------------------------------------------------------------------------- -@pytest.mark.gpu +@pytest.mark.cuda class TestMixedDevices: def test_inconsistent_warns_and_moves_to_first(self): if not torch.cuda.is_available(): diff --git a/tests/unit/utils/test_triton_dispatch.py b/tests/unit/utils/test_triton_dispatch.py deleted file mode 100644 index 20f07697..00000000 --- a/tests/unit/utils/test_triton_dispatch.py +++ /dev/null @@ -1,86 +0,0 @@ -"""Unit tests for the shared Triton/eager dispatch engine (CPU-only).""" - -import pytest -import torch - -from torchref.utils import ( - Engine, - get_engine, - set_engine, - should_use_triton, - triton_available, - use_engine, -) - -pytestmark = pytest.mark.unit - - -def test_engine_roundtrip(): - prev = get_engine() - try: - set_engine(Engine.TRITON) - assert get_engine() is Engine.TRITON - set_engine(Engine.EAGER) - assert get_engine() is Engine.EAGER - finally: - set_engine(prev) - - -def test_use_engine_nests_and_restores(): - prev = get_engine() - try: - set_engine(Engine.AUTO) - with use_engine(Engine.EAGER): - assert get_engine() is Engine.EAGER - with use_engine(Engine.TRITON): - assert get_engine() is Engine.TRITON - assert get_engine() is Engine.EAGER - assert get_engine() is Engine.AUTO - finally: - set_engine(prev) - - -def test_use_engine_restores_on_exception(): - prev = get_engine() - try: - set_engine(Engine.AUTO) - with pytest.raises(ValueError): - with use_engine(Engine.EAGER): - raise ValueError("boom") - assert get_engine() is Engine.AUTO - finally: - set_engine(prev) - - -def test_eager_engine_never_uses_triton(): - cpu = torch.zeros(3) - assert should_use_triton(cpu, engine=Engine.EAGER) is False - # even with a (hypothetical) cuda fp32 tensor, EAGER forces eager - assert should_use_triton(None, engine=Engine.EAGER) is False - - -def test_auto_false_on_cpu_and_fp64(): - assert should_use_triton(torch.zeros(3, dtype=torch.float32), engine=Engine.AUTO) is False - assert should_use_triton(torch.zeros(3, dtype=torch.float64), engine=Engine.AUTO) is False - - -def test_triton_engine_raises_on_cpu_or_fp64(): - with pytest.raises(RuntimeError): - should_use_triton(torch.zeros(3, dtype=torch.float32), engine=Engine.TRITON) - with pytest.raises(RuntimeError): - should_use_triton(torch.zeros(3, dtype=torch.float64), engine=Engine.TRITON) - - -def test_none_tensors_are_ignored(): - # all-None probes: capability is vacuously satisfied; AUTO depends only on - # triton availability (False on a CPU-only box, True if triton importable). - assert should_use_triton(None, None, engine=Engine.AUTO) is triton_available() - - -def test_should_use_triton_reads_global_engine_by_default(): - prev = get_engine() - try: - set_engine(Engine.EAGER) - assert should_use_triton(torch.zeros(3)) is False - finally: - set_engine(prev) diff --git a/torchref/__init__.py b/torchref/__init__.py index 092bfef9..2ff89bbd 100644 --- a/torchref/__init__.py +++ b/torchref/__init__.py @@ -51,7 +51,7 @@ General utilities and debugging tools. """ -__version__ = "0.6.1" +__version__ = "0.6.2" import os diff --git a/torchref/base/alignment/normalization.py b/torchref/base/alignment/normalization.py index 2599a47e..e51e600b 100644 --- a/torchref/base/alignment/normalization.py +++ b/torchref/base/alignment/normalization.py @@ -14,7 +14,7 @@ import torch from torchref.base.math_torch import U_to_matrix -from torchref.config import get_default_device +from torchref.config import normalize_device def compute_radial_shells( @@ -46,8 +46,7 @@ def compute_radial_shells( shell_centers : torch.Tensor Shell centers in Angstroms^-1, shape (n_shells,). """ - if device is None: - device = get_default_device() + device = normalize_device(device) s_min = 1.0 / d_max # Low resolution end s_max = 1.0 / d_min # High resolution end diff --git a/torchref/base/alignment/rotation.py b/torchref/base/alignment/rotation.py index a07f8251..ad0d4399 100644 --- a/torchref/base/alignment/rotation.py +++ b/torchref/base/alignment/rotation.py @@ -8,7 +8,7 @@ import numpy as np import torch -from torchref.config import dtypes, get_default_device, get_float_dtype +from torchref.config import dtypes, get_float_dtype, normalize_device def rotate_coords_torch(coords, phi, rho): @@ -261,8 +261,7 @@ def random_rotation_uniform( torch.Tensor Rotation matrices with shape (n, 3, 3) or (3, 3) if n=1. """ - if device is None: - device = get_default_device() + device = normalize_device(device) if dtype is None: dtype = get_float_dtype() # Sample uniform random numbers diff --git a/torchref/base/direct_summation/__init__.py b/torchref/base/direct_summation/__init__.py index 3df4873e..9993b8ed 100644 --- a/torchref/base/direct_summation/__init__.py +++ b/torchref/base/direct_summation/__init__.py @@ -65,7 +65,7 @@ def compute_scattering_factors_batch( # Capability-based backend dispatch (Triton on CUDA+fp32, else checkpointed # eager). Keep ``triton_ds`` itself lazy (loaded inside dispatch) so a broken # Triton install never breaks ``import torchref``. -from .dispatch import Engine, ds_aniso, ds_iso +from .dispatch import ds_aniso, ds_iso __all__ = [ # Scattering factor batch helper @@ -82,7 +82,6 @@ def compute_scattering_factors_batch( "core_deformation", "multiplication_quasi_complex_tensor", # Dispatch - "Engine", "ds_iso", "ds_aniso", ] diff --git a/torchref/base/direct_summation/_backends.py b/torchref/base/direct_summation/_backends.py new file mode 100644 index 00000000..e77720b5 --- /dev/null +++ b/torchref/base/direct_summation/_backends.py @@ -0,0 +1,125 @@ +"""The direct-summation dispatch policy, as one table. + +Companion to ``electron_density/_backends.py``; see :mod:`torchref.utils.backends` for what +each field means. + +Two things about this table are worth reading before changing it. + +**There is no Metal direct-summation kernel.** The Metal shader is a *density splat* only, so +DS on MPS is ``_checkpointed_*`` running on-device -- a real production path, and one the +table states by simply having no MPS row. + +**The float32 requirement here is policy, not capability**, and the policy is mild. The +Triton kernel casts every input to float32 itself (``triton_ds._cols_f32``), so it would +happily consume float64; the gate declines instead, so that a float64 configuration is not +silently served at float32 precision. That is a consistency rule, not a rescue: measured on +a 40-atom / 60-reflection scene, downcasting the real-valued inputs moves ``F`` by 7.2e-06 +relative to ``mean|F|``, three orders of magnitude inside the 1e-2 amplitude tolerance this +package works to. + +It matters even less than that suggests, because the interesting case was never reachable. +Under a float64 configuration ``xyz_frac`` is float64, and the previous gate probed +``xyz_frac`` -- so it already declined. The only thing widening the probe set catches is +hand-mixed dtypes (float32 coordinates beside a float64 ``s`` or ``adp``), which no caller in +the repo produces. Treat this as making the gate check what its kernel actually consumes, +consistent with the density table, rather than as a fix for anything observed. + +``hkl`` is deliberately **not** probed, and probing it would be wrong. Miller indices are +integers -- ``int32`` from the MTZ reader, and integer-valued even in a float dtype, since +symmetry maps ``h -> h.R`` with integer ``R``. The f64->f32 round-trip is bit-exact for any +\\|h\\| < 2**24 (real structures reach a few hundred), and measurably so: downcasting ``hkl`` +alone changes ``F`` by exactly 0.0. Probing it would only manufacture a false negative. + +The dtype rule exempts integer tensors for the same reason -- without that, the ``int32`` +production dtype would read as a capability failure and disable the kernel outright. +""" + +from __future__ import annotations + +import torch + +from torchref.utils.backends import Backend, BackendTable +from torchref.utils.backends import triton_available + +_THIS = "torchref.base.direct_summation._backends" + + +def why_unavailable(): + """``None`` if the direct-summation Triton kernels can run, else why they cannot. + + Two separable questions, and both have to be asked. ``triton_available()`` answers "is + the ``triton`` package importable"; this additionally answers "does ``triton_ds`` + itself import", which can fail on its own -- the module evaluates + ``tl.constexpr(...)`` at import time and that is version-sensitive. The density + backend has the same split. + """ + if not triton_available(): + return "triton is not importable" + try: + from torchref.base.direct_summation import triton_ds # noqa: F401 + except Exception as exc: # noqa: BLE001 - any import failure is a reason, not a crash + return ( + "the direct-summation Triton kernels could not be imported " + f"({type(exc).__name__}: {exc})" + ) + return None + + +# --------------------------------------------------------------------------- +# Signature adapters +# --------------------------------------------------------------------------- +# The two backends do not share a signature: the checkpointed path takes a trailing +# ``max_memory_gb`` reflection-chunk budget and the Triton kernel has nothing to bound -- +# it forms only a ``(BLOCK_H, N)`` tile in registers. Rather than add a parameter the +# kernel ignores (a lie in a public signature, and the exact drift just removed from the +# Metal wrappers), the budget is dropped in a named adapter. Named, not a lambda, so it +# appears in a traceback and can be patched. +def _ds_iso_triton(hkl, s, xyz_frac, occ, adp, A, B, max_memory_gb): + """Isotropic Triton DS; ``max_memory_gb`` is not applicable and is dropped here.""" + from torchref.base.direct_summation.triton_ds import ds_iso_triton + + return ds_iso_triton(hkl, s, xyz_frac, occ, adp, A, B) + + +def _ds_aniso_triton(hkl, s_vec, xyz_frac, occ, U, A, B, max_memory_gb): + """Anisotropic Triton DS; ``max_memory_gb`` is not applicable and is dropped here.""" + from torchref.base.direct_summation.triton_ds import ds_aniso_triton + + return ds_aniso_triton(hkl, s_vec, xyz_frac, occ, U, A, B) + + +DS_BACKENDS = BackendTable( + name="direct-summation structure factors", + backends=( + Backend( + name="ds_triton", + kernel=(_THIS, "_ds_iso_triton", "_ds_aniso_triton"), + device="cuda", + dtypes=(torch.float32,), + # Every argument except ``hkl`` (position 0), whose dtype provably costs + # nothing -- see the module docstring. + probes=(1, 2, 3, 4, 5, 6), + probe=(_THIS, "why_unavailable"), + expect_available="cuda", + on_failure="degrade", + second_order=False, + ), + Backend( + name="checkpointed", + kernel=( + "torchref.base.direct_summation.dispatch", + "_checkpointed_iso", + "_checkpointed_aniso", + ), + # METAL is here because there is no Metal DS kernel; see the module docstring. + # TRITON is absent, which is what makes that engine strict. + expect_available="always", + on_failure="raise", + # First-order only: the backward replays each chunk under ``enable_grad`` but + # without ``create_graph``, so a second derivative raises rather than returning + # something wrong. ``_eager_*`` is the double-differentiable reference and is + # deliberately not in this table -- it is not a production dispatch target. + second_order=False, + ), + ), +) diff --git a/torchref/base/direct_summation/dispatch.py b/torchref/base/direct_summation/dispatch.py index 18a862fc..d3005af8 100644 --- a/torchref/base/direct_summation/dispatch.py +++ b/torchref/base/direct_summation/dispatch.py @@ -1,15 +1,13 @@ """Backend dispatch for direct-summation structure factors. -Selection uses the shared, capability-based :class:`~torchref.utils.Engine` -(see :mod:`torchref.utils.triton_dispatch`): for a given ``(device, dtype)`` -there is exactly one best path, so the backend is *derived* rather than -configured. +Which backend runs is read from +:data:`torchref.base.direct_summation._backends.DS_BACKENDS` -- device, dtype, which +device and dtype, how availability is probed, and what a failure means all live there, so +this module holds the maths and the entry points but no selection logic. The mechanics are in +:mod:`torchref.utils.backends`. -- CUDA + float32 + Triton available -> custom Triton kernels (``triton_ds``) -- everything else (CPU, MPS, non-float32, no Triton) -> checkpointed eager - -An explicit ``Engine`` override (per-call ``engine=`` or the process-wide -``use_engine``/``set_engine``) forces a path for tests and benchmarks. +``force_portable`` (per call, or scoped with ``use_portable()``) pins the portable reference +path; otherwise the fastest usable backend is chosen automatically. All backends compute a **P1** structure-factor sum only. Crystallographic symmetry is applied outside, in :meth:`SfDS.compute_structure_factors`. @@ -30,37 +28,12 @@ _estimate_batch_size_aniso, aniso_structure_factor_torched, ) -from torchref.utils.triton_dispatch import Engine, should_use_triton +from torchref.utils.backends import run_or_degrade, select TWO_PI = 2.0 * math.pi NEG_TWO_PI_SQ = -2.0 * (math.pi**2) -# --------------------------------------------------------------------------- -# Lazy Triton probe (mirrors electron_density/main.py:_get_*_triton) -# --------------------------------------------------------------------------- -_triton_iso_fn = None -_triton_aniso_fn = None -_triton_checked = False - - -def _load_triton(): - """Import the Triton kernels once; leave the fns ``None`` if unavailable.""" - global _triton_iso_fn, _triton_aniso_fn, _triton_checked - if not _triton_checked: - try: - from torchref.base.direct_summation.triton_ds import ( - ds_iso_triton, - ds_aniso_triton, - ) - - _triton_iso_fn = ds_iso_triton - _triton_aniso_fn = ds_aniso_triton - except Exception: - pass - _triton_checked = True - - # --------------------------------------------------------------------------- # P1 identity symmetry (lives here, not in sf_ds) # --------------------------------------------------------------------------- @@ -247,40 +220,40 @@ def _eager_aniso(hkl, s_vec, xyz_frac, occ, U, A, B, max_memory_gb): # --------------------------------------------------------------------------- # Dispatch entry points # --------------------------------------------------------------------------- -# The coarse Triton-vs-eager gate is the shared ``should_use_triton`` (device, -# dtype, engine). The DS-kernel-module availability is then checked here: under -# AUTO a missing/failing kernel falls back to checkpointed eager; under -# Engine.TRITON it raises. -def ds_iso(hkl, s, xyz_frac, occ, adp, A, B, *, engine=Engine.AUTO, max_memory_gb=None): - """Isotropic P1 structure factors via the selected backend.""" +def _dispatch(aniso, hkl, geom, xyz_frac, occ, third, A, B, force_portable, max_memory_gb): + """Select a backend from ``DS_BACKENDS`` and run it. + + Imported inside the call rather than at module scope: the table names this module as a + kernel source, so a module-level import would close a cycle. + """ + from torchref.base.direct_summation._backends import DS_BACKENDS + + args = (hkl, geom, xyz_frac, occ, third, A, B) + backend = select(DS_BACKENDS, args, force_portable=force_portable) + return run_or_degrade(DS_BACKENDS, backend, aniso, *args, max_memory_gb) + + +def ds_iso(hkl, s, xyz_frac, occ, adp, A, B, *, force_portable=None, max_memory_gb=None): + """Isotropic P1 structure factors via the selected backend. + + ``force_portable=None`` defers to the process-wide setting, so + ``with use_portable(): ...`` reaches this call. Passing ``True`` pins the checkpointed + reference path regardless. + """ if xyz_frac.shape[0] == 0: return torch.zeros(hkl.shape[0], dtype=torch.complex64, device=hkl.device) - if should_use_triton(xyz_frac, engine=engine): - _load_triton() - if _triton_iso_fn is not None: - try: - return _triton_iso_fn(hkl, s, xyz_frac, occ, adp, A, B) - except Exception: - if engine is Engine.TRITON: - raise - # AUTO: fall back to checkpointed eager - elif engine is Engine.TRITON: - raise RuntimeError("Direct-summation Triton kernels are unavailable") - return _checkpointed_iso(hkl, s, xyz_frac, occ, adp, A, B, max_memory_gb) - - -def ds_aniso(hkl, s_vec, xyz_frac, occ, U, A, B, *, engine=Engine.AUTO, max_memory_gb=None): - """Anisotropic P1 structure factors via the selected backend.""" + return _dispatch( + False, hkl, s, xyz_frac, occ, adp, A, B, force_portable, max_memory_gb + ) + + +def ds_aniso(hkl, s_vec, xyz_frac, occ, U, A, B, *, force_portable=None, max_memory_gb=None): + """Anisotropic P1 structure factors via the selected backend. + + See :func:`ds_iso` on ``force_portable=None``. + """ if xyz_frac.shape[0] == 0: return torch.zeros(hkl.shape[0], dtype=torch.complex64, device=hkl.device) - if should_use_triton(xyz_frac, engine=engine): - _load_triton() - if _triton_aniso_fn is not None: - try: - return _triton_aniso_fn(hkl, s_vec, xyz_frac, occ, U, A, B) - except Exception: - if engine is Engine.TRITON: - raise - elif engine is Engine.TRITON: - raise RuntimeError("Direct-summation Triton kernels are unavailable") - return _checkpointed_aniso(hkl, s_vec, xyz_frac, occ, U, A, B, max_memory_gb) + return _dispatch( + True, hkl, s_vec, xyz_frac, occ, U, A, B, force_portable, max_memory_gb + ) diff --git a/torchref/base/electron_density/__init__.py b/torchref/base/electron_density/__init__.py index bc23b981..ddd3a62d 100644 --- a/torchref/base/electron_density/__init__.py +++ b/torchref/base/electron_density/__init__.py @@ -8,7 +8,7 @@ - Scatter operations for map building The headline entry point is :func:`build_electron_density` (from ``main``): the -central, ``Engine``-dispatched per-atom variable-radius density builder that +central, table-dispatched per-atom variable-radius density builder that takes atomic parameters and returns the full density map. """ diff --git a/torchref/base/electron_density/_backends.py b/torchref/base/electron_density/_backends.py new file mode 100644 index 00000000..eed753e9 --- /dev/null +++ b/torchref/base/electron_density/_backends.py @@ -0,0 +1,101 @@ +"""The electron-density splat dispatch policy, as one table. + +Every criterion for choosing a density kernel is a field in a row below: which device, +which dtypes, how availability is probed, whether a runtime failure +may degrade, and whether the kernel composes to second order. Reading this file answers +"which kernel runs for MPS + float64?" without tracing an if/elif ladder +through three modules. + +All four wrappers share one signature -- +``(density_map, xyz, adp_or_u, occ, A, B, inv_frac, frac, radius_per_atom)`` -- which is +what lets the dispatch site be a single call rather than a per-kernel adapter. + +Ordering within the table is not load-bearing. The three non-base backends are pairwise +device-disjoint (CUDA / CPU / MPS), so at most one can ever match; the base case matches +everything and is what selection falls through to. + +See :mod:`torchref.utils.backends` for what each field means and why it exists. +""" + +from __future__ import annotations + +import torch + +from torchref.utils.backends import Backend, BackendTable + +_CUDA = "torchref.base.electron_density.kernels.cuda.variable_radius" +_MPS = "torchref.base.electron_density.kernels.mps.variable_radius" +_SPHERE = "torchref.base.electron_density.kernels.cpu.sphere_splat" +_PORTABLE = "torchref.base.electron_density.kernels.cpu.variable_radius" + +#: Argument positions carrying the device/dtype contract: +#: ``density_map, xyz, adp_or_u, occ, A, B``. The trailing three -- the two cell matrices +#: and the per-atom radius -- are excluded because no gate has ever probed them, and +#: widening the contract to cover them would make two currently-working backends stricter +#: for no demonstrated bug. They *are* read as raw pointers by the C++ and CUDA kernels, so +#: that remains an open question rather than a settled one. +_ATOM_ARGS = (0, 1, 2, 3, 4, 5) + +DENSITY_BACKENDS = BackendTable( + name="electron-density splat", + backends=( + Backend( + name="cuda_triton", + kernel=(_CUDA, "add_isotropic_cuda_var", "add_anisotropic_cuda_var"), + device="cuda", + dtypes=(torch.float32,), + probes=_ATOM_ARGS, + probe=(_CUDA, "why_unavailable"), + expect_available="cuda", + # A failed launch on an available GPU is a capability miss worth degrading + # from; the kernel clones the density map before accumulating, so the + # fallback cannot double-count. + on_failure="degrade", + second_order=False, + ), + Backend( + name="mps_metal", + kernel=(_MPS, "add_isotropic_mps_var", "add_anisotropic_mps_var"), + device="mps", + dtypes=(torch.float32,), + probes=_ATOM_ARGS, + probe=("torchref.base.electron_density.kernels.mps.compile", + "why_unavailable"), + expect_available="mps", + on_failure="degrade", + second_order=False, + ), + Backend( + name="cpu_sphere", + kernel=(_SPHERE, "add_isotropic_cpu_sphere_var", + "add_anisotropic_cpu_sphere_var"), + device="cpu", + dtypes=(torch.float32, torch.float64), + # Uniformity, not membership: the kernel picks one ``scalar_t`` from the output + # map and then reads every other tensor through a raw pointer of that type, so a + # float64 map beside float32 atoms would be a 2x out-of-bounds read. + require_uniform_dtype=True, + probes=_ATOM_ARGS, + probe=(_SPHERE, "why_unavailable"), + # A pure C++ build with no hardware requirement, so it must work everywhere. The + # expectation is what makes a broken build a CI failure instead of a skip. + expect_available="always", + # Deliberately not "degrade". The extension having failed to *build* already + # yields a reason from the probe, so this governs only a runtime throw from a + # kernel that compiled -- a genuine bug. Falling back would turn wrong results + # into a ~100x slowdown whose numbers still look plausible, because the portable + # splat implements the same truncation contract. + on_failure="raise", + second_order=True, + ), + Backend( + name="portable", + kernel=(_PORTABLE, "add_isotropic_plain_var", "add_anisotropic_plain_var"), + # The base case: no device or dtype restriction, so it always matches. That is + # what makes ``select`` unable to fail. + expect_available="always", + on_failure="raise", + second_order=True, + ), + ), +) diff --git a/torchref/base/electron_density/kernels/__init__.py b/torchref/base/electron_density/kernels/__init__.py index 9bf5df45..77a5bbde 100644 --- a/torchref/base/electron_density/kernels/__init__.py +++ b/torchref/base/electron_density/kernels/__init__.py @@ -2,11 +2,9 @@ Optimized density-splatting kernels, organized by device. Layout: -- ``cpu/`` — the per-atom variable-radius CPU splats - (``variable_radius.py``: the production CPU AUTO grouped-separable and the - portable plain-scatter splats), the shared separable density core - (``separable.py``), the aniso splat, the C++ parallel scatter - (``scatter.py`` / ``scatter_dispatch.py``), and the JIT reference. +- ``cpu/`` — the production fused C++ spherical-cutoff splat + (``sphere_splat.py``), the portable plain-scatter splats + (``variable_radius.py``), and the JIT reference. - ``cuda/`` — the production variable-radius work-queue kernels (``variable_radius.py``: ``WorkQueueGridDensity{,Aniso}``) plus the legacy fixed-radius fused Triton kernel (``fused.py``, benchmark-only). @@ -28,13 +26,6 @@ clear_cache, ) -# Triton kernels are optional (require the triton package / CUDA). -try: - from .cuda.fused import fused_add_to_map_gpu - _HAS_TRITON = True -except ImportError: - _HAS_TRITON = False - __all__ = [ "vectorized_add_to_map", "build_electron_density", @@ -43,5 +34,24 @@ "warmup", "get_cache_dir", "clear_cache", - "fused_add_to_map_gpu", ] + +# Triton kernels are optional (they require the triton package and a GPU). +# +# ``except Exception``, not ``except ImportError``: this runs during ``import torchref``, +# and a Triton install that is present but broken -- a driver or LLVM version skew, the +# common real-world failure -- raises something other than ImportError on import. Catching +# only ImportError meant such a host could not import torchref at all, even though +# ``torchref.utils.triton_available()`` was written to absorb exactly this and would have +# reported False. The sibling guard in ``cuda/variable_radius.py`` already used the wider +# clause. +try: + from .cuda.fused import fused_add_to_map_gpu + + # Appended rather than listed unconditionally. ``fused_add_to_map_gpu`` is bound only + # if the import succeeded, so naming it in a static ``__all__`` made + # ``from torchref.base.electron_density.kernels import *`` raise AttributeError on any + # host without Triton. + __all__.append("fused_add_to_map_gpu") +except Exception: # pragma: no cover - depends on the host's triton install + pass diff --git a/torchref/base/electron_density/kernels/cpu/_cpp_build.py b/torchref/base/electron_density/kernels/cpu/_cpp_build.py new file mode 100644 index 00000000..d039ad07 --- /dev/null +++ b/torchref/base/electron_density/kernels/cpu/_cpp_build.py @@ -0,0 +1,133 @@ +"""Shared ``load_inline`` harness for the CPU C++ kernels. + +Extracted from ``scatter.py`` so the fused sphere splat (``sphere_splat.py``) does +not carry a second copy of it. Everything here is cluster-deployment plumbing that +must not drift between the two extensions: + +* **POSIX ``lockf`` locking** instead of PyTorch's ``FileBaton``. FileBaton is a + file-existence lock, so a process killed mid-compile leaves it behind and blocks + every future import. ``fcntl.lockf`` record locks are enforced by the filesystem + (so they work across NFS/GPFS nodes) and released by the kernel on process death, + even SIGKILL. +* **Per-microarchitecture build directory**, keyed on the CPU model. Without it a + ``-march=native`` binary built on one cluster node raises Illegal Instruction on + a node with a different CPU (e.g. AMD vs Intel). +* **ninja on PATH** — pip-installed ninja lives next to ``sys.executable``, which + is not on PATH on compute nodes. +* **GCC >= 9** via ``/opt/rh/gcc-toolset-*``, required by PyTorch C++ extensions. +* **No ``-fopenmp`` on macOS** — Apple Clang rejects it (no bundled OpenMP + runtime), so kernels must provide a ``std::thread`` fallback under + ``#ifndef _OPENMP``. + +Failures are returned, never raised: every caller degrades to a slower pure-torch +path, so a missing compiler is a performance problem and not an outage. +""" + +from __future__ import annotations + +import os +import sys +import traceback +from typing import Optional, Tuple + +from torch.utils.cpp_extension import load_inline + + +def cpu_tag() -> str: + """A tag identifying this CPU's microarchitecture, for the build directory.""" + import platform + + tag = platform.machine() + try: + with open("/proc/cpuinfo") as f: + for line in f: + if line.startswith("model name"): + # e.g. "EPYC_7443P" or "Xeon_Gold_6248" + return line.split(":")[1].strip().replace(" ", "_") + except OSError: + pass + return tag + + +def build_extension( + name: str, + cpp_source: str, + extra_cflags: Optional[list] = None, +) -> Tuple[object, Optional[Tuple[str, str]]]: + """Compile ``cpp_source`` into an extension module. + + Parameters + ---------- + name : str + Extension name; also the ``PYBIND11_MODULE`` name the source must use. + cpp_source : str + Complete translation unit, including its own ``PYBIND11_MODULE`` block. + extra_cflags : list of str, optional + Appended after the shared ``-O3 -march=native``. ``-fopenmp`` is added + automatically on every platform except macOS. + + Returns + ------- + (module, error) + ``(module, None)`` on success, ``(None, (message, traceback))`` on any + failure. Never raises. + """ + try: + import fcntl + except ImportError: + # Non-POSIX (Windows): no record locks, so don't attempt a build at all. + return None, ("fcntl unavailable (non-POSIX platform)", "") + + try: + # pip-installed ninja sits next to the interpreter, not on PATH. + bin_dir = os.path.dirname(sys.executable) + if bin_dir not in os.environ.get("PATH", ""): + os.environ["PATH"] = bin_dir + ":" + os.environ.get("PATH", "") + # PyTorch C++ extensions need GCC >= 9. + for toolset in ("14", "13", "12"): + gcc = f"/opt/rh/gcc-toolset-{toolset}/root/usr/bin/g++" + if os.path.isfile(gcc): + os.environ["CXX"] = gcc + os.environ["CC"] = gcc.replace("g++", "gcc") + break + + build_dir = os.path.join( + os.environ.get( + "TORCH_EXTENSIONS_DIR", + os.path.join(os.path.expanduser("~"), ".cache", "torch_extensions"), + ), + f"{name}_{cpu_tag()}", + ) + os.makedirs(build_dir, exist_ok=True) + + cflags = ["-O3", "-march=native"] + list(extra_cflags or []) + ldflags: list = [] + if sys.platform != "darwin": # Apple Clang has no bundled OpenMP runtime + cflags.append("-fopenmp") + ldflags.append("-fopenmp") + + lock_fd = os.open( + os.path.join(build_dir, "compile.lock"), os.O_CREAT | os.O_RDWR + ) + try: + fcntl.lockf(lock_fd, fcntl.LOCK_EX) + # Clear a stale PyTorch FileBaton lock left by a killed process. + try: + os.unlink(os.path.join(build_dir, "lock")) + except FileNotFoundError: + pass + module = load_inline( + name=name, + cpp_sources=[cpp_source], + extra_cflags=cflags, + extra_ldflags=ldflags, + build_directory=build_dir, + verbose=False, + ) + finally: + fcntl.lockf(lock_fd, fcntl.LOCK_UN) + os.close(lock_fd) + except Exception as e: # noqa: BLE001 - any build failure degrades gracefully + return None, (f"{type(e).__name__}: {e}", traceback.format_exc()) + + return module, None diff --git a/torchref/base/electron_density/kernels/cpu/aniso.py b/torchref/base/electron_density/kernels/cpu/aniso.py deleted file mode 100644 index 7dbb8793..00000000 --- a/torchref/base/electron_density/kernels/cpu/aniso.py +++ /dev/null @@ -1,63 +0,0 @@ -"""Anisotropic Gaussian box-splat for CPU / MPS (``Engine.AUTO``). - -Mirrors the isotropic separable box-splat (center index -> local voxel cube -> -per-axis ``d_frac`` -> structured scatter) but evaluates the full 3D anisotropic -Gaussian over the cube (the cross-terms U12/U13/U23 prevent 1D factorization). -""" - -import math - -import torch - -from torchref.config import dtypes - - -def _aniso_density_cube(d_frac, frac_matrix, Minv, A_norm): - """Anisotropic Gaussian density over the local voxel cube. - - The aniso analogue of ``_separable_density`` — but the 3D Gaussian - does NOT factorize across axes (cross-terms U12/U13/U23), so it builds the - full Cartesian offset cube and evaluates the quadratic form directly. The - component loop keeps peak memory at O(C*n^3). - - Parameters - ---------- - d_frac : (C, 3, n) — PBC-wrapped fractional offsets per axis. - frac_matrix : (3, 3) — fractional -> Cartesian. - Minv : (C, 5, 3, 3) — inverse of M_g = (B_g*I + 8*pi^2*U)/4. - A_norm : (C, 5) — A * occ * pi^1.5 / sqrt(det M_g). - - Returns - ------- - (C, n, n, n) density cube. - """ - pi_sq = math.pi * math.pi - da = d_frac[:, 0, :][:, :, None, None] # (C, n, 1, 1) - db = d_frac[:, 1, :][:, None, :, None] # (C, 1, n, 1) - dc = d_frac[:, 2, :][:, None, None, :] # (C, 1, 1, n) - fm = frac_matrix - # Cartesian offset cube r = frac_matrix @ d_frac (each (C, n, n, n)) - cx = fm[0, 0] * da + fm[0, 1] * db + fm[0, 2] * dc - cy = fm[1, 0] * da + fm[1, 1] * db + fm[1, 2] * dc - cz = fm[2, 0] * da + fm[2, 1] * db + fm[2, 2] * dc - - C = d_frac.shape[0] - n = d_frac.shape[2] - density_cube = d_frac.new_zeros(C, n, n, n) - for g in range(Minv.shape[1]): - m00 = Minv[:, g, 0, 0][:, None, None, None] - m11 = Minv[:, g, 1, 1][:, None, None, None] - m22 = Minv[:, g, 2, 2][:, None, None, None] - m01 = Minv[:, g, 0, 1][:, None, None, None] - m02 = Minv[:, g, 0, 2][:, None, None, None] - m12 = Minv[:, g, 1, 2][:, None, None, None] - q = ( - m00 * cx * cx - + m11 * cy * cy - + m22 * cz * cz - + 2.0 * (m01 * cx * cy + m02 * cx * cz + m12 * cy * cz) - ) - density_cube = density_cube + A_norm[:, g, None, None, None] * torch.exp( - -pi_sq * q - ) - return density_cube diff --git a/torchref/base/electron_density/kernels/cpu/jit_reference.py b/torchref/base/electron_density/kernels/cpu/jit_reference.py index f96b9a86..5de72590 100644 --- a/torchref/base/electron_density/kernels/cpu/jit_reference.py +++ b/torchref/base/electron_density/kernels/cpu/jit_reference.py @@ -8,7 +8,7 @@ Architecture: - CPU: JIT-scripted kernel using einsum with metric tensor (efficient for CPU) -- GPU: when the shared ``Engine`` permits Triton (CUDA + float32, engine +- GPU: when the shared targets gate permits Triton (CUDA + float32, dispatch AUTO/TRITON), ``vectorized_add_to_map`` selects the Triton fused branch (``fused_add_to_map_gpu``) first; otherwise it falls back to the pure-torch, double-differentiable ``_add_to_map_gpu_simple`` (JIT-scripted batch matmul). @@ -29,7 +29,7 @@ import os import torch -from torchref.utils.triton_dispatch import should_use_triton +from torchref.base.targets._dispatch import use_triton # ============================================================================= # Cache directory for JIT kernels @@ -379,10 +379,10 @@ def vectorized_add_to_map( """ Add atoms to density map using ITC92 Gaussian parameterization. - Backend is chosen by the shared ``Engine`` via ``should_use_triton``: on + Backend is chosen by the shared targets gate via ``use_triton``: on GPU it uses the Triton fused kernel when Triton is permitted (CUDA+float32, engine AUTO/TRITON) and the pure-torch ``_add_to_map_gpu_simple`` otherwise - (Engine.EAGER, float64, or Triton unavailable). CPU uses the JIT kernel. + (force_portable, float64, or Triton unavailable). CPU uses the JIT kernel. Parameters ---------- @@ -417,11 +417,11 @@ def vectorized_add_to_map( input unchanged. Callers must always use the returned value. """ if density_map.device.type == "cuda": - # The shared Engine is the only switch: use the Triton kernel when it - # permits (CUDA + float32, engine AUTO/TRITON); otherwise — Engine.EAGER, + # The shared targets gate is the only switch: use the Triton kernel when it + # permits (CUDA + float32); otherwise — force_portable, # float64, or Triton unavailable — the pure-torch, double-differentiable # ``_add_to_map_gpu_simple``. - if should_use_triton(xyz): + if use_triton(xyz): triton_fn = _get_triton_kernel() if triton_fn is not None: return triton_fn( @@ -489,7 +489,7 @@ def build_electron_density( This alias takes the *voxel-level* signature (precomputed ``surrounding_coords`` / ``voxel_indices``) and is distinct from the top-level :func:`torchref.base.electron_density.main.build_electron_density`, - which takes *atomic* parameters and performs the full ``Engine``-based + which takes *atomic* parameters and performs the full table-based variable-radius dispatch. Parameters diff --git a/torchref/base/electron_density/kernels/cpu/scatter.py b/torchref/base/electron_density/kernels/cpu/scatter.py deleted file mode 100644 index 0016a5ce..00000000 --- a/torchref/base/electron_density/kernels/cpu/scatter.py +++ /dev/null @@ -1,647 +0,0 @@ -"""Parallel partitioned scatter_add for structured (wa + wbwc) indices. - -This kernel is float32-only by design: the C++ kernel hardcodes -``torch::kFloat32`` for the density/output buffers. The Python wrapper -auto-casts non-float32 inputs to float32 and emits a one-time warning per -input dtype — so callers running with ``dtypes.float = torch.float64`` keep -working, just at float32 precision through this kernel. - -Uses a C++ kernel compiled via torch.utils.cpp_extension.load_inline. -Each thread owns a contiguous output partition and only accumulates scatter -elements that fall in its range. With atoms sorted by 1D center, per-thread -early-exit skips ~(T-1)/T of atoms — giving near-linear scaling. - -Threading backend is selected at compile time: - * OpenMP on Linux (and any compiler that accepts -fopenmp) - * std::thread fallback otherwise (notably Apple Clang on macOS, whose - -fopenmp is rejected because the OpenMP runtime is not bundled) - -Two index dtypes are exposed: - * int32 — default fast path. Fits any realistic crystallographic grid - (nx*ny*nz < 2**31, roughly 1300^3 voxels) and halves the - index bandwidth of the inner loop. - * int64 — fallback for the rare edge cases that exceed INT32_MAX - (e.g. very large cryo-EM grids). -The Python wrapper picks the binding from wa.dtype. - -Autograd: forward is the custom C++ scatter, backward is a standard gather -(embarrassingly parallel, no custom kernel needed). -""" - -import warnings - -import torch -from torch.utils.cpp_extension import load_inline - -_WARNED_CAST_DTYPES: set = set() - -_CPP_SRC = r""" -#include -#include - -#ifdef _OPENMP -#include -#else -#include -#include -#endif - -// Parallel partitioned scatter_add from structured indices. -// -// Output space is divided among threads. Each thread only writes to its own -// [lo, hi) segment — zero synchronization. -// Atoms sorted by 1D center give efficient early-exit per thread. -// -// IdxT is int32_t (default fast path) or int64_t (for huge grids). -template -torch::Tensor structured_scatter_add_impl( - torch::Tensor output, - torch::Tensor wa, - torch::Tensor wbwc, - torch::Tensor values, - int64_t nx, int64_t ny, int64_t nz) -{ - TORCH_CHECK(output.is_contiguous() && output.scalar_type() == torch::kFloat32); - TORCH_CHECK(values.is_contiguous() && values.scalar_type() == torch::kFloat32); - TORCH_CHECK(wa.is_contiguous()); - TORCH_CHECK(wbwc.is_contiguous()); - - const int64_t M = output.size(0); - const int64_t C = wa.size(0); - const int64_t ny_nz = ny * nz; - const int64_t nxyz = nx * ny_nz; - - float* __restrict__ out_p = output.data_ptr(); - const IdxT* __restrict__ wa_p = wa.data_ptr(); - const IdxT* __restrict__ wbwc_p = wbwc.data_ptr(); - const float* __restrict__ val_p = values.data_ptr(); - - // Global wbwc bounds (for atom-level early exit) - IdxT wbwc_min = wbwc_p[0], wbwc_max = wbwc_p[0]; - for (int64_t i = 1; i < C * ny_nz; i++) { - IdxT v = wbwc_p[i]; - if (v < wbwc_min) wbwc_min = v; - if (v > wbwc_max) wbwc_max = v; - } - - auto worker = [&](int64_t lo, int64_t hi) { - for (int64_t c = 0; c < C; c++) { - const IdxT* wa_row = wa_p + c * nx; - - // Atom-level early exit: check wa range vs partition - IdxT amin = wa_row[0], amax = wa_row[0]; - for (int64_t i = 1; i < nx; i++) { - IdxT v = wa_row[i]; - if (v < amin) amin = v; - if (v > amax) amax = v; - } - if ((int64_t)amax + wbwc_max < lo || (int64_t)amin + wbwc_min >= hi) continue; - - const IdxT* wbwc_row = wbwc_p + c * ny_nz; - const float* val_base = val_p + c * nxyz; - - for (int64_t ix = 0; ix < nx; ix++) { - IdxT wa_val = wa_row[ix]; - // x-offset early exit - if ((int64_t)wa_val + wbwc_max < lo || (int64_t)wa_val + wbwc_min >= hi) continue; - - const float* val_ix = val_base + ix * ny_nz; - - for (int64_t iyz = 0; iyz < ny_nz; iyz++) { - int64_t idx = (int64_t)wa_val + (int64_t)wbwc_row[iyz]; - if (idx >= lo && idx < hi) { - out_p[idx] += val_ix[iyz]; - } - } - } - } - }; - -#ifdef _OPENMP - #pragma omp parallel - { - int tid = omp_get_thread_num(); - int nth = omp_get_num_threads(); - int64_t lo = (int64_t)tid * M / nth; - int64_t hi = (int64_t)(tid + 1) * M / nth; - worker(lo, hi); - } -#else - int nth = (int)std::thread::hardware_concurrency(); - if (nth < 1) nth = 1; - std::vector threads; - threads.reserve(nth); - for (int tid = 0; tid < nth; tid++) { - int64_t lo = (int64_t)tid * M / nth; - int64_t hi = (int64_t)(tid + 1) * M / nth; - threads.emplace_back(worker, lo, hi); - } - for (auto& t : threads) t.join(); -#endif - - return output; -} - -// Parallel structured gather (backward of scatter_add). -// -// grad_cube[c, ix, iyz] = grad_output[wa[c,ix] + wbwc[c,iyz]] -// -// Each atom's output is independent — parallelize over atoms directly. -// No index tensor allocation, no fancy indexing overhead. -// Note: the worker_gather lambda below is used only by the non-OpenMP branch; -// under _OPENMP the same loop body is duplicated inline in the parallel-for. -template -torch::Tensor structured_gather_impl( - torch::Tensor grad_output, - torch::Tensor wa, - torch::Tensor wbwc, - int64_t nx, int64_t ny, int64_t nz) -{ - TORCH_CHECK(grad_output.is_contiguous() && grad_output.scalar_type() == torch::kFloat32); - TORCH_CHECK(wa.is_contiguous()); - TORCH_CHECK(wbwc.is_contiguous()); - - const int64_t C = wa.size(0); - const int64_t ny_nz = ny * nz; - const int64_t nxyz = nx * ny_nz; - - auto grad_cube = torch::empty({C, nxyz}, grad_output.options()); - - const float* __restrict__ go_p = grad_output.data_ptr(); - const IdxT* __restrict__ wa_p = wa.data_ptr(); - const IdxT* __restrict__ wbwc_p = wbwc.data_ptr(); - float* __restrict__ gc_p = grad_cube.data_ptr(); - - auto worker_gather = [&](int64_t c_lo, int64_t c_hi) { - for (int64_t c = c_lo; c < c_hi; c++) { - const IdxT* wa_row = wa_p + c * nx; - const IdxT* wbwc_row = wbwc_p + c * ny_nz; - float* out_row = gc_p + c * nxyz; - - for (int64_t ix = 0; ix < nx; ix++) { - IdxT wa_val = wa_row[ix]; - float* dst = out_row + ix * ny_nz; - - for (int64_t iyz = 0; iyz < ny_nz; iyz++) { - dst[iyz] = go_p[(int64_t)wa_val + (int64_t)wbwc_row[iyz]]; - } - } - } - }; - -#ifdef _OPENMP - #pragma omp parallel for schedule(static) - for (int64_t c = 0; c < C; c++) { - const IdxT* wa_row = wa_p + c * nx; - const IdxT* wbwc_row = wbwc_p + c * ny_nz; - float* out_row = gc_p + c * nxyz; - - for (int64_t ix = 0; ix < nx; ix++) { - IdxT wa_val = wa_row[ix]; - float* dst = out_row + ix * ny_nz; - - for (int64_t iyz = 0; iyz < ny_nz; iyz++) { - dst[iyz] = go_p[(int64_t)wa_val + (int64_t)wbwc_row[iyz]]; - } - } - } -#else - int nth = (int)std::thread::hardware_concurrency(); - if (nth < 1) nth = 1; - std::vector threads; - threads.reserve(nth); - for (int tid = 0; tid < nth; tid++) { - int64_t c_lo = (int64_t)tid * C / nth; - int64_t c_hi = (int64_t)(tid + 1) * C / nth; - threads.emplace_back(worker_gather, c_lo, c_hi); - } - for (auto& t : threads) t.join(); -#endif - - return grad_cube; -} - -// Wrappers that enforce the matching dtype, so a wrong-dtype tensor gets a -// clear error instead of being reinterpreted by data_ptr(). -torch::Tensor structured_scatter_add_i32( - torch::Tensor output, torch::Tensor wa, torch::Tensor wbwc, torch::Tensor values, - int64_t nx, int64_t ny, int64_t nz) -{ - TORCH_CHECK(wa.scalar_type() == torch::kInt32, "wa must be int32"); - TORCH_CHECK(wbwc.scalar_type() == torch::kInt32, "wbwc must be int32"); - return structured_scatter_add_impl(output, wa, wbwc, values, nx, ny, nz); -} -torch::Tensor structured_scatter_add_i64( - torch::Tensor output, torch::Tensor wa, torch::Tensor wbwc, torch::Tensor values, - int64_t nx, int64_t ny, int64_t nz) -{ - TORCH_CHECK(wa.scalar_type() == torch::kInt64, "wa must be int64"); - TORCH_CHECK(wbwc.scalar_type() == torch::kInt64, "wbwc must be int64"); - return structured_scatter_add_impl(output, wa, wbwc, values, nx, ny, nz); -} -torch::Tensor structured_gather_i32( - torch::Tensor grad_output, torch::Tensor wa, torch::Tensor wbwc, - int64_t nx, int64_t ny, int64_t nz) -{ - TORCH_CHECK(wa.scalar_type() == torch::kInt32, "wa must be int32"); - TORCH_CHECK(wbwc.scalar_type() == torch::kInt32, "wbwc must be int32"); - return structured_gather_impl(grad_output, wa, wbwc, nx, ny, nz); -} -torch::Tensor structured_gather_i64( - torch::Tensor grad_output, torch::Tensor wa, torch::Tensor wbwc, - int64_t nx, int64_t ny, int64_t nz) -{ - TORCH_CHECK(wa.scalar_type() == torch::kInt64, "wa must be int64"); - TORCH_CHECK(wbwc.scalar_type() == torch::kInt64, "wbwc must be int64"); - return structured_gather_impl(grad_output, wa, wbwc, nx, ny, nz); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("structured_scatter_add_i32", &structured_scatter_add_i32, - "Partitioned scatter_add with int32 structured (wa, wbwc) indices"); - m.def("structured_scatter_add_i64", &structured_scatter_add_i64, - "Partitioned scatter_add with int64 structured (wa, wbwc) indices"); - m.def("structured_gather_i32", &structured_gather_i32, - "Structured gather (backward) with int32 indices"); - m.def("structured_gather_i64", &structured_gather_i64, - "Structured gather (backward) with int64 indices"); -} -""" - -# --------------------------------------------------------------------------- -# Lazy compilation with POSIX lockf locking (works across cluster nodes). -# -# PyTorch's load_inline uses FileBaton (file-existence lock) which stays -# behind when a process is killed mid-compile, blocking all future imports. -# We wrap it with fcntl.lockf (POSIX record locks) which: -# - are enforced by the filesystem → work across NFS/GPFS cluster nodes -# - are released by the kernel on process death (even SIGKILL) -# First process compiles; all others (same node or different) reuse cache. -# --------------------------------------------------------------------------- -_module = None -_module_failed = False -# Captured diagnostic from the most recent failed compile attempt. Populated -# by _get_module() on failure so test/debug code can surface the underlying -# error instead of just "module not available". Format: (error_str, traceback_str). -_module_error: "tuple[str, str] | None" = None - - -def _get_module(): - global _module, _module_failed, _module_error - if _module is not None: - return _module - if _module_failed: - return None - - import os - import sys - import traceback - - try: - import fcntl - except ImportError: - # fcntl is not available on non-POSIX platforms (e.g. Windows) - _module_failed = True - _module_error = ("fcntl unavailable (non-POSIX platform)", "") - return None - - try: - # Ensure ninja (installed via pip) is on PATH for compute nodes - bin_dir = os.path.dirname(sys.executable) - if bin_dir not in os.environ.get("PATH", ""): - os.environ["PATH"] = bin_dir + ":" + os.environ.get("PATH", "") - # Need GCC >= 9 for PyTorch C++ extensions - for toolset in ("14", "13", "12"): - gcc = f"/opt/rh/gcc-toolset-{toolset}/root/usr/bin/g++" - if os.path.isfile(gcc): - os.environ["CXX"] = gcc - os.environ["CC"] = gcc.replace("g++", "gcc") - break - - # Per-microarchitecture build directory — prevents Illegal Instruction - # when different cluster nodes have different CPUs (e.g., AMD vs Intel). - import platform - - cpu_tag = platform.machine() - try: - # Use the CPU model to distinguish microarchitectures - with open("/proc/cpuinfo") as f: - for line in f: - if line.startswith("model name"): - # e.g. "EPYC_7443P" or "Xeon_Gold_6248" - cpu_tag = line.split(":")[1].strip().replace(" ", "_") - break - except OSError: - pass - build_dir = os.path.join( - os.environ.get( - "TORCH_EXTENSIONS_DIR", - os.path.join(os.path.expanduser("~"), ".cache", "torch_extensions"), - ), - f"cpu_scatter_{cpu_tag}", - ) - os.makedirs(build_dir, exist_ok=True) - - # Apple Clang on macOS rejects -fopenmp (no bundled OpenMP runtime). - # Kernel falls back to std::thread via #ifndef _OPENMP — no libomp, - # no Homebrew, no external dependency required. - is_apple_clang = sys.platform == "darwin" - extra_cflags = ["-O3", "-march=native"] - extra_ldflags: list[str] = [] - if not is_apple_clang: - extra_cflags.append("-fopenmp") - extra_ldflags.append("-fopenmp") - - # fcntl.lockf uses POSIX record locks (fcntl F_SETLKW) which are: - # 1. filesystem-level → work across NFS/GPFS cluster nodes - # 2. released by kernel on process death, even SIGKILL - lock_fd = os.open( - os.path.join(build_dir, "compile.lock"), os.O_CREAT | os.O_RDWR - ) - try: - fcntl.lockf(lock_fd, fcntl.LOCK_EX) - # Clear any stale PyTorch FileBaton lock from a killed process - try: - os.unlink(os.path.join(build_dir, "lock")) - except FileNotFoundError: - pass - _module = load_inline( - name="cpu_scatter", - cpp_sources=[_CPP_SRC], - extra_cflags=extra_cflags, - extra_ldflags=extra_ldflags, - build_directory=build_dir, - verbose=False, - ) - finally: - fcntl.lockf(lock_fd, fcntl.LOCK_UN) - os.close(lock_fd) - except Exception as e: - _module_failed = True - _module_error = (f"{type(e).__name__}: {e}", traceback.format_exc()) - return None - - return _module - - -# --------------------------------------------------------------------------- -# Autograd wrapper -# --------------------------------------------------------------------------- - - -class _StructuredScatterAdd(torch.autograd.Function): - """scatter_add with structured (wa, wbwc) indices. - - Forward: C++ partitioned scatter (parallel, no conflicts) - Backward: standard gather (embarrassingly parallel) - """ - - @staticmethod - def forward(ctx, density_cube, wa, wbwc, map_size): - # density_cube: (C, nx, ny, nz) — requires grad, must be float32 - # wa: (C, nx) — int32 or int64, no grad - # wbwc: (C, ny, nz) — same dtype as wa - if density_cube.dtype != torch.float32: - if density_cube.dtype not in _WARNED_CAST_DTYPES: - warnings.warn( - f"cpu_scatter is float32-only; casting density_cube from " - f"{density_cube.dtype} to float32 (precision will be reduced).", - stacklevel=3, - ) - _WARNED_CAST_DTYPES.add(density_cube.dtype) - density_cube = density_cube.to(torch.float32) - if wa.dtype != wbwc.dtype: - raise TypeError( - f"wa.dtype ({wa.dtype}) and wbwc.dtype ({wbwc.dtype}) must match" - ) - if wa.dtype == torch.int32: - if map_size > _INT32_MAX: - raise RuntimeError( - f"map_size {map_size} exceeds INT32_MAX ({_INT32_MAX}); " - "pass int64 indices for grids this large." - ) - scatter_fn_name = "structured_scatter_add_i32" - gather_fn_name = "structured_gather_i32" - elif wa.dtype == torch.int64: - scatter_fn_name = "structured_scatter_add_i64" - gather_fn_name = "structured_gather_i64" - else: - raise TypeError(f"wa.dtype must be int32 or int64, got {wa.dtype}") - - C, nx, ny, nz = density_cube.shape - ctx.save_for_backward(wa, wbwc) - ctx.cube_shape = density_cube.shape - ctx.gather_fn_name = gather_fn_name - - mod = _get_module() - if mod is None: - err = _module_error[0] if _module_error else "unknown reason" - raise RuntimeError( - f"C++ cpu_scatter module not available ({err}). " - "See torchref.base.electron_density.kernels.cpu.scatter._module_error for the full traceback." - ) - result = torch.zeros( - map_size, dtype=density_cube.dtype, device=density_cube.device - ) - getattr(mod, scatter_fn_name)( - result, - wa.contiguous(), - wbwc.reshape(C, ny * nz).contiguous(), - density_cube.reshape(C, nx * ny * nz).contiguous(), - nx, - ny, - nz, - ) - return result - - @staticmethod - def backward(ctx, grad_output): - wa, wbwc = ctx.saved_tensors - C, nx, ny, nz = ctx.cube_shape - # The scatter is linear, so its VJP is the adjoint gather. Routing - # through the ``_StructuredGather`` Function (rather than calling the - # C++ gather imperatively) keeps the operation in the autograd graph, - # so higher-order derivatives compose: gather's own backward is the - # scatter again (the two are mutual adjoints). Without this, a C++ - # gather returns a graph-less tensor and ``create_graph=True`` silently - # drops the second-order term through the scatter transpose. - grad_cube = _StructuredGather.apply(grad_output, wa, wbwc, ctx.cube_shape) - return grad_cube, None, None, None - - -class _StructuredGather(torch.autograd.Function): - """Adjoint of :class:`_StructuredScatterAdd` (structured gather). - - ``gather(grid)[c] = grid[scatter_index(c)]``. Linear, so its VJP is the - scatter. Pairing the two as mutual adjoints makes both double-(and higher-) - differentiable with no Hessian math. - """ - - @staticmethod - def forward(ctx, grid, wa, wbwc, cube_shape): - if grid.dtype != torch.float32: - if grid.dtype not in _WARNED_CAST_DTYPES: - warnings.warn( - f"cpu_scatter is float32-only; casting from {grid.dtype} to " - f"float32 (precision will be reduced).", - stacklevel=3, - ) - _WARNED_CAST_DTYPES.add(grid.dtype) - grid = grid.to(torch.float32) - C, nx, ny, nz = cube_shape - mod = _get_module() - if mod is None: - err = _module_error[0] if _module_error else "unknown reason" - raise RuntimeError( - f"C++ cpu_scatter module not available ({err}). " - "See torchref.base.electron_density.kernels.cpu.scatter._module_error for the full traceback." - ) - # int32 / int64 gather fn name matches the scatter's index dtype. - gather_fn_name = ( - "structured_gather_i32" - if wa.dtype == torch.int32 - else "structured_gather_i64" - ) - grad_cube = getattr(mod, gather_fn_name)( - grid.contiguous(), - wa.contiguous(), - wbwc.reshape(C, ny * nz).contiguous(), - nx, - ny, - nz, - ).reshape(C, nx, ny, nz) - ctx.save_for_backward(wa, wbwc) - ctx.map_size = int(grid.numel()) - return grad_cube - - @staticmethod - def backward(ctx, grad_cube): - wa, wbwc = ctx.saved_tensors - # Adjoint of the gather is the scatter — reuse the differentiable - # Function so a third-order graph would also compose. - grad_grid = _StructuredScatterAdd.apply(grad_cube, wa, wbwc, ctx.map_size) - return grad_grid, None, None, None - - -_INT32_MAX = 2**31 - 1 - - -def structured_scatter_add(density_cube, wa, wbwc, map_size): - """Differentiable parallel scatter_add using structured indices. - - Dispatches to the int32 or int64 kernel based on ``wa.dtype``. int32 is - the default fast path (halves index bandwidth); int64 is available for - grids larger than INT32_MAX voxels. - - Parameters - ---------- - density_cube : Tensor (C, nx, ny, nz) float32 - Values to scatter (from _separable_density). - wa : Tensor (C, nx) int32 or int64 - Precomputed x-axis scatter indices. - wbwc : Tensor (C, ny, nz) same dtype as wa - Precomputed yz-plane scatter indices. - map_size : int - Total number of voxels in flat density map. For int32 indices, - must be <= INT32_MAX. - - Returns - ------- - Tensor (map_size,) float32 - Scattered result. Differentiable w.r.t. density_cube. - """ - return _StructuredScatterAdd.apply(density_cube, wa, wbwc, map_size) - - -def structured_scatter_add_inplace(out_flat, density_cube, wa, wbwc): - """Low-level in-place structured scatter: ``out_flat[idx] += cube`` (no autograd). - - The raw building block for :class:`_ScatterAccumulate`. The C++ kernel already - does ``out[idx] += ...``, so passing ``out_flat`` straight in as the output - buffer accumulates in place — no fresh ``zeros(map_size)`` and no out-of-place - ``out + result`` add (the two full-grid touches per chunk that the functional - :func:`structured_scatter_add` incurs). - - Returns ``out_flat`` (mutated), or ``None`` if the C++ module is unavailable - (caller should fall back to the functional path). - """ - mod = _get_module() - if mod is None: - return None - if density_cube.dtype != torch.float32: - density_cube = density_cube.to(torch.float32) - C, nx, ny, nz = density_cube.shape - if wa.dtype == torch.int32: - if out_flat.numel() > _INT32_MAX: - raise RuntimeError( - f"map_size {out_flat.numel()} exceeds INT32_MAX ({_INT32_MAX}); " - "pass int64 indices for grids this large." - ) - fn_name = "structured_scatter_add_i32" - elif wa.dtype == torch.int64: - fn_name = "structured_scatter_add_i64" - else: - raise TypeError(f"wa.dtype must be int32 or int64, got {wa.dtype}") - getattr(mod, fn_name)( - out_flat, - wa.contiguous(), - wbwc.reshape(C, ny * nz).contiguous(), - density_cube.reshape(C, nx * ny * nz).contiguous(), - nx, - ny, - nz, - ) - return out_flat - - -class _ScatterAccumulate(torch.autograd.Function): - """Differentiable in-place accumulating scatter: ``out += scatter(cube)``. - - ``out_flat`` is mutated in place (one shared buffer accumulated across chunks), - so there is NO per-chunk ``zeros(map_size)`` and NO per-chunk out-of-place - ``out + result`` add — the two full-grid touches that dominate scatter-bound, - many-chunk structures in the functional path. - - Correctness under autograd: the scatter is linear, so for ``out = out_in + - scatter(cube)`` the VJP is ``grad_out`` w.r.t. ``out_in`` (identity) and - ``gather(grad_out)`` w.r.t. ``cube``. Crucially the backward needs ONLY the - (un-mutated) indices ``wa``/``wbwc`` and ``grad_out`` — never the mutated - buffer — so the in-place version bump invalidates no saved tensor and the - chunk-to-chunk in-place chain differentiates correctly. The gather is routed - through :class:`_StructuredGather` so higher-order graphs still compose. - """ - - @staticmethod - def forward(ctx, out_flat, density_cube, wa, wbwc): - if density_cube.dtype != torch.float32: - if density_cube.dtype not in _WARNED_CAST_DTYPES: - warnings.warn( - f"cpu_scatter is float32-only; casting density_cube from " - f"{density_cube.dtype} to float32 (precision will be reduced).", - stacklevel=3, - ) - _WARNED_CAST_DTYPES.add(density_cube.dtype) - density_cube = density_cube.to(torch.float32) - ctx.save_for_backward(wa, wbwc) - ctx.cube_shape = density_cube.shape - structured_scatter_add_inplace(out_flat, density_cube, wa, wbwc) - ctx.mark_dirty(out_flat) - return out_flat - - @staticmethod - def backward(ctx, grad_out): - wa, wbwc = ctx.saved_tensors - # d out / d out_in = identity -> grad_out passes through unchanged. - # d out / d cube = adjoint gather of grad_out. - grad_cube = _StructuredGather.apply(grad_out, wa, wbwc, ctx.cube_shape) - return grad_out, grad_cube, None, None - - -def structured_scatter_accumulate(out_flat, density_cube, wa, wbwc): - """Differentiable ``out_flat += scatter(density_cube)``, accumulated in place. - - Returns the (mutated) ``out_flat``. Differentiable w.r.t. both ``out_flat`` and - ``density_cube``. Use in the chunk loop to accumulate every chunk into one - shared buffer with no per-chunk full-grid allocation or add. - """ - return _ScatterAccumulate.apply(out_flat, density_cube, wa, wbwc) diff --git a/torchref/base/electron_density/kernels/cpu/scatter_dispatch.py b/torchref/base/electron_density/kernels/cpu/scatter_dispatch.py deleted file mode 100644 index 69750697..00000000 --- a/torchref/base/electron_density/kernels/cpu/scatter_dispatch.py +++ /dev/null @@ -1,80 +0,0 @@ -"""Device-aware structured scatter for the box-splat density kernels. - -On CPU, dispatches to the custom C++ parallel scatter (``cpu_scatter``) when the -extension built successfully; otherwise (and on every non-CPU device) falls back -to PyTorch's ``scatter_add_``. -""" - -import torch - -# Lazy-loaded C++ parallel scatter for CPU -_cpp_scatter_fn = None -_cpp_scatter_accumulate_fn = None -_cpp_scatter_checked = False - - -def _get_cpp_scatter(): - """Return the C++ parallel scatter_add, or None if unavailable. - - Eagerly triggers the C++ compilation so that failures (missing ninja, - unsupported compiler flags, etc.) are caught here rather than mid-calculation. - Also caches the in-place variant used by the no-grad fast path. - """ - global _cpp_scatter_fn, _cpp_scatter_accumulate_fn, _cpp_scatter_checked - if not _cpp_scatter_checked: - try: - from torchref.base.electron_density.kernels.cpu.scatter import ( - _get_module, - structured_scatter_accumulate, - structured_scatter_add, - ) - - # Trigger compilation now — _get_module returns None on failure - if _get_module() is not None: - _cpp_scatter_fn = structured_scatter_add - _cpp_scatter_accumulate_fn = structured_scatter_accumulate - except Exception: - pass - _cpp_scatter_checked = True - return _cpp_scatter_fn - - -def _do_structured_scatter( - density_cube: torch.Tensor, - wa: torch.Tensor, - wbwc: torch.Tensor, - density_flat: torch.Tensor, - map_size: int, -) -> torch.Tensor: - """Pick the fastest structured scatter for the device. - - On CPU, dispatches to the custom C++ kernel (``cpu_scatter``) when - available — partitioned, no atomics, ~2× faster than PyTorch's stock - ``scatter_add_``. On every other device (MPS, CUDA, CPU without the - extension built) falls back to PyTorch ``scatter_add_`` with int64 - indices. - - Returns the resulting flat density tensor. The C++ path accumulates - out-of-place; the ``scatter_add_`` fallback mutates ``density_flat`` - in place. Both return the up-to-date tensor. - """ - if density_cube.device.type == "cpu": - cpp_fn = _get_cpp_scatter() - if cpp_fn is not None: - # Differentiable in-place accumulate: density_flat += scatter(cube), - # mutating one shared buffer across chunks. Avoids the per-chunk - # zeros(map_size) + out-of-place full-grid add (two full-grid touches - # per chunk) that the functional path incurs and that dominate - # scatter-bound, many-chunk structures. Gradients are preserved - # (the scatter is linear; see _ScatterAccumulate). - if _cpp_scatter_accumulate_fn is not None: - return _cpp_scatter_accumulate_fn(density_flat, density_cube, wa, wbwc) - return density_flat + cpp_fn(density_cube, wa, wbwc, map_size) - # Fallback: PyTorch scatter_add_ requires int64 indices. - idx_flat = wa[:, :, None, None] + wbwc[:, None, :, :] - density_flat.scatter_add_( - 0, - idx_flat.reshape(-1).to(torch.int64), - density_cube.reshape(-1), - ) - return density_flat diff --git a/torchref/base/electron_density/kernels/cpu/separable.py b/torchref/base/electron_density/kernels/cpu/separable.py deleted file mode 100644 index 99ee87d5..00000000 --- a/torchref/base/electron_density/kernels/cpu/separable.py +++ /dev/null @@ -1,155 +0,0 @@ -"""Separable Gaussian box-splat for isotropic atoms (CPU / MPS shared core). - -Factorizes exp(-alpha * r^T G r) into 1D Gaussians per axis with 2D cross-term -corrections for non-orthogonal cells, keeping peak memory low. ``_separable_density`` -is the shared core reused by the variable-radius CPU splat -(``cpu/variable_radius.py::add_isotropic_cpu_separable_var``). -""" - -import math - -import torch - -from torchref.config import dtypes - -# Peak-element budget for the triclinic general path's batched (C, N_comp, n^3) -# exp intermediate. Atoms are sub-batched so this is the memory ceiling -# regardless of the caller's chunk size. 16M floats ~= 64 MB. -_GENERAL_BUDGET = 16_000_000 - - -def _separable_density( - d_frac: torch.Tensor, - alpha: torch.Tensor, - A_norm: torch.Tensor, - G: torch.Tensor, - has_ab: bool, - has_ac: bool, - has_bc: bool, -) -> torch.Tensor: - """Separable Gaussian density evaluation. - - Factorizes exp(-alpha * r^T G r) into 1D Gaussians per axis with 2D - cross-term corrections for non-orthogonal cells. Batches all corrections - across the 5 ITC92 components and uses einsum where possible. - - For non-orthogonal cells, cross-term exponents are combined with the - relevant diagonal exponents before taking exp() to avoid float32 overflow - (exp(-big) * exp(+big) = 0 * inf = NaN). Each combined 2D block - exponent corresponds to a principal sub-matrix of G (positive definite), - guaranteeing the exponent is always <= 0 and exp() is in (0, 1]. - - Dispatch by crystal system for optimal performance: - - Orthogonal (no cross terms): separable 1D products + einsum - - Hexagonal (ab only): combined ab exponent + einsum with e_c - - Monoclinic (ac only): combined ac exponent + einsum with e_b - - General (bc, or multiple cross terms): full 3D exponent per component - - Parameters - ---------- - d_frac : (C, 3, n_axis) — fractional distances per axis, PBC-wrapped. - alpha : (C, N_comp) — pi^2 / B_total. - A_norm : (C, N_comp) — weighted amplitudes. - G : (3, 3) — metric tensor frac_matrix.T @ frac_matrix. - has_ab : bool — whether G[0,1] cross-term is significant. - has_ac : bool — whether G[0,2] cross-term is significant. - has_bc : bool — whether G[1,2] cross-term is significant. - - Returns - ------- - (C, n_axis, n_axis, n_axis) density cube. - """ - # --- Convert fractional → Cartesian per-axis --- - cell_lengths = torch.sqrt(torch.diagonal(G)) - d_cart = d_frac * cell_lengths[None, :, None] # (C, 3, n) - - # --- 1D exponents (always <= 0) --- - da2 = d_cart[:, 0, :] ** 2 - db2 = d_cart[:, 1, :] ** 2 - dc2 = d_cart[:, 2, :] ** 2 - - log_a = -alpha.unsqueeze(2) * da2.unsqueeze(1) # (C, Nc, n) - log_b = -alpha.unsqueeze(2) * db2.unsqueeze(1) - log_c = -alpha.unsqueeze(2) * dc2.unsqueeze(1) - - if not (has_ab or has_ac or has_bc): - # ---- Orthogonal cells: pure separable, all exp() args <= 0 ---- - e_a = torch.exp(log_a) - e_b = torch.exp(log_b) - e_c = torch.exp(log_c) - e_ab = e_a.unsqueeze(3) * e_b.unsqueeze(2) - return torch.einsum("cg,cgij,cgk->cijk", A_norm, e_ab, e_c) - - # --- Cross-term coefficients --- - cos_gamma = G[0, 1] / (cell_lengths[0] * cell_lengths[1]) - cos_beta = G[0, 2] / (cell_lengths[0] * cell_lengths[2]) - cos_alpha = G[1, 2] / (cell_lengths[1] * cell_lengths[2]) - - da = d_cart[:, 0, :] - db = d_cart[:, 1, :] - dc = d_cart[:, 2, :] - alpha_4d = alpha[:, :, None, None] # (C, Nc, 1, 1) - - if has_ab and not has_ac and not has_bc: - # ---- Hexagonal / trigonal: only ab cross-term ---- - # Combined 2D exponent: -alpha*(da2 + db2 + 2*cos_gamma*da*db) - # = -alpha * d_ab^T G_ab d_ab <= 0 (G_ab positive definite) - prod_ab = da.unsqueeze(2) * db.unsqueeze(1) - log_ab = ( - log_a[:, :, :, None] - + log_b[:, :, None, :] - + (-2.0 * alpha_4d * cos_gamma * prod_ab[:, None, :, :]) - ) - slice_ab = torch.exp(log_ab) # (C, Nc, n, n), all in (0, 1] - e_c = torch.exp(log_c) - return torch.einsum("cg,cgij,cgk->cijk", A_norm, slice_ab, e_c) - - if has_ac and not has_ab and not has_bc: - # ---- Monoclinic (beta != 90): only ac cross-term ---- - # Combined 2D exponent: -alpha*(da2 + dc2 + 2*cos_beta*da*dc) - # = -alpha * d_ac^T G_ac d_ac <= 0 (G_ac positive definite) - prod_ac = da.unsqueeze(2) * dc.unsqueeze(1) - log_ac = ( - log_a[:, :, :, None] - + log_c[:, :, None, :] - + (-2.0 * alpha_4d * cos_beta * prod_ac[:, None, :, :]) - ) - e_ac = torch.exp(log_ac) # (C, Nc, n_a, n_c), all in (0, 1] - e_b = torch.exp(log_b) - return torch.einsum("cg,cgj,cgik->cijk", A_norm, e_b, e_ac) - - # ---- General path (triclinic, or multiple cross-terms) ---- - # The per-component exponent is -alpha[:, g] * Q, where the quadratic form - # Q = da^2 + db^2 + dc^2 + 2*cos_gamma*da*db + 2*cos_beta*da*dc + 2*cos_alpha*db*dc - # is the cell geometry (r^T G r in cell-length-normalized units) and is - # INDEPENDENT of the ITC92 component g (alpha[:, g] factors out entirely). - # So build Q ONCE (it is positive semidefinite -> -alpha*Q <= 0, no overflow), - # then evaluate exp(-alpha[:, g] * Q) for all components in a SINGLE batched - # exp + einsum rather than a per-component loop with repeated geometry rebuilds. - # The batched (C, N_comp, n, n, n) intermediate is 5x the cube, so atoms are - # processed in sub-batches sized to keep peak memory under _GENERAL_BUDGET - # floats (the result is identical to the loop, just fewer/larger exp calls -- - # ~1.5-2x faster on the triclinic path which is dominated by the exp). - C = d_frac.shape[0] - n = d_frac.shape[2] - Nc = alpha.shape[1] - sub = max(1, _GENERAL_BUDGET // max(1, Nc * n * n * n)) - parts = [] - for s in range(0, C, sub): - e = min(s + sub, C) - Q = ( - da2[s:e, :, None, None] - + db2[s:e, None, :, None] - + dc2[s:e, None, None, :] - ) # (c, n, n, n) — broadcasting materializes a full contiguous cube - if has_ab: - Q = Q + (2.0 * cos_gamma * da[s:e].unsqueeze(2) * db[s:e].unsqueeze(1)).unsqueeze(3) - if has_ac: - Q = Q + (2.0 * cos_beta * da[s:e].unsqueeze(2) * dc[s:e].unsqueeze(1)).unsqueeze(2) - if has_bc: - Q = Q + (2.0 * cos_alpha * db[s:e].unsqueeze(2) * dc[s:e].unsqueeze(1)).unsqueeze(1) - # (c, N_comp, n, n, n): one exp over all components, then weighted sum. - comp = torch.exp(-alpha[s:e, :, None, None, None] * Q.unsqueeze(1)) - parts.append(torch.einsum("cg,cgijk->cijk", A_norm[s:e], comp)) - - return parts[0] if len(parts) == 1 else torch.cat(parts, dim=0) diff --git a/torchref/base/electron_density/kernels/cpu/sphere_splat.py b/torchref/base/electron_density/kernels/cpu/sphere_splat.py new file mode 100644 index 00000000..e54dbe58 --- /dev/null +++ b/torchref/base/electron_density/kernels/cpu/sphere_splat.py @@ -0,0 +1,801 @@ +"""Fused per-atom spherical-cutoff density splat for CPU (C++), with autograd. + +This is the production CPU path for :func:`build_electron_density`. It is a +transliteration of the Metal kernels in ``kernels/mps/_shaders.py`` -- deliberately +so: the point is that CPU, CUDA and Metal implement *one* truncation contract, so +``torchref.sigma_cutoff_ed`` means the same thing on every device. + +The canonical contract, shared with the Triton and Metal kernels: + + voxel v receives atom i's full 5-Gaussian density iff ||w||^2 <= r_i^2, + +where ``w`` is the minimum-image **Cartesian atom->voxel** vector (so the sphere is +centred on the atom, not on its nearest grid node) and ``r_i`` is the raw +``radius_policy`` radius. Enumeration bound is the triclinic-correct per-axis +half-width ``ceil(r_i * n_axis * ||inv_frac row_axis||)``. There is no +grid-dependent requantization of the radius and no diagonal-metric approximation. + +Why fused rather than the older grouped-separable splat +------------------------------------------------------ +The separable kernel this replaced (since deleted) factorized the Gaussian into 1D +per-axis exponentials, which needs a *uniform box per launch* -- hence a cube cutoff and +all the bucket-by-box-size machinery. Measured on 4000 atoms / 2.3M voxels / 0.40 A +sampling, the cube touches 19.3M voxels where the sphere touches 8.3M, and +materializing the ``(C,L,L,L)`` intermediate cube dominates the runtime: this fused +kernel is ~1.5x faster than the separable one even using ``std::exp``, and ~2.5x +faster with ``fast_exp``. So the spherical cutoff is not a performance concession +here -- it is both the correct geometry and the faster one. + +Threading +--------- +Forward partitions the **output** by x-plane via ``at::parallel_for``: each thread +owns a disjoint set of planes, writes only into them, and needs no atomics (a +cheap wrapped-interval test rejects atoms that miss the partition). Backward +partitions over **atoms**, so each thread owns its atom's gradient slots +exclusively -- also no atomics, exactly as in the Metal backward. + +Precision +--------- +float32 uses ``fast_exp`` (branchless 2^x bit-trick plus a degree-5 minimax +polynomial), the CPU analogue of the ``metal::fast::exp`` the Metal kernels already +use. Measured agreement with ``std::exp`` is rel L2 1.9e-5 -- some 40x below the +7.9e-4 amplitude-truncation floor at the default 3 sigma. float64 uses ``std::exp``, +since a float64 caller is precision-motivated by definition. + +Gradients flow to ``xyz``, ``adp``/``u`` and ``occ``, with identity to the incoming +``density_map``; ``A``/``B`` and the cell matrices get none -- the same set as the +CUDA and Metal kernels. Backward is first-order only; double backward must use +``force_portable`` (the plain splat in ``variable_radius.py``). + +Unlike the CUDA and Metal entry points this kernel takes no per-Gaussian +``coeff_mask``: that argument is all-ones at every existing call site, so it is +omitted rather than allocated. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +from torchref.base.electron_density.kernels.cpu._cpp_build import build_extension + +_CPP_SRC = r""" +#include +#include +#include +#include +#include + +// --------------------------------------------------------------------------- +// exp: fast branchless variant for float, libm for double. +// +// exp(x) = 2^(x*log2(e)): the integer part goes straight into the IEEE exponent +// field, the fractional part through a degree-5 minimax polynomial on [0,1). +// No branches, so the innermost voxel loop stays vectorizable. This mirrors +// metal::fast::exp, which the Metal kernels already use. +// --------------------------------------------------------------------------- +static inline float fast_exp(float x) { + x = x < -87.0f ? -87.0f : x; // below this, exp underflows + const float t = x * 1.44269504088896341f; // log2(e) + const float n = std::floor(t); + const float f = t - n; + const float p = 1.0f + f * (0.6931471805f + f * (0.2402265069f + + f * (0.0555041087f + f * (0.0096181291f + f * 0.0013333558f)))); + const int32_t bits = (int32_t)((n + 127.0f) * 8388608.0f) & 0x7f800000; + float scale; + std::memcpy(&scale, &bits, sizeof(scale)); + return p * scale; +} + +// Overloads rather than explicit specializations: the templated kernels below +// call kexp() unqualified and overload resolution picks by scalar_t. +static inline float kexp(float x) { return fast_exp(x); } +static inline double kexp(double x) { return std::exp(x); } + +template struct K { + static constexpr scalar_t PI_1P5 = (scalar_t)5.568327996831708; // pi^1.5 + static constexpr scalar_t PI_SQ = (scalar_t)9.869604401089358; // pi^2 + static constexpr scalar_t TWO_PI_SQ= (scalar_t)19.739208802178716; // 2*pi^2 +}; + +// Cell geometry + this atom's anchor, shared by every kernel below. +template +struct Cell { + scalar_t uax,uay,uaz, ubx,uby,ubz, ucx,ucy,ucz; // Cartesian voxel steps + scalar_t inva,invb,invc; // inv_frac row norms + scalar_t Auc; // |u_c|^2 + scalar_t fnx,fny,fnz; + int nx,ny,nz; +}; + +template +static Cell make_cell(const scalar_t* fm, const scalar_t* im, + int nx, int ny, int nz) { + Cell c; + c.nx=nx; c.ny=ny; c.nz=nz; + c.fnx=(scalar_t)nx; c.fny=(scalar_t)ny; c.fnz=(scalar_t)nz; + // frac columns are the cell vectors a,b,c; divided by n gives the voxel step. + c.uax=fm[0]/c.fnx; c.uay=fm[3]/c.fnx; c.uaz=fm[6]/c.fnx; + c.ubx=fm[1]/c.fny; c.uby=fm[4]/c.fny; c.ubz=fm[7]/c.fny; + c.ucx=fm[2]/c.fnz; c.ucy=fm[5]/c.fnz; c.ucz=fm[8]/c.fnz; + c.inva=std::sqrt(im[0]*im[0]+im[1]*im[1]+im[2]*im[2]); + c.invb=std::sqrt(im[3]*im[3]+im[4]*im[4]+im[5]*im[5]); + c.invc=std::sqrt(im[6]*im[6]+im[7]*im[7]+im[8]*im[8]); + c.Auc=c.ucx*c.ucx+c.ucy*c.ucy+c.ucz*c.ucz; + return c; +} + +template +struct Anchor { + int cix,ciy,ciz, bhx,bhy,bhz; + scalar_t w0x,w0y,w0z, rc2; +}; + +template +static inline Anchor make_anchor( + const Cell& c, const scalar_t* xyz, const scalar_t* fm, + const scalar_t* im, int64_t a, scalar_t rc2) +{ + Anchor g; + g.rc2 = rc2; + const scalar_t r = std::sqrt(rc2); + // Per-axis bounding box of the Cartesian r-sphere in index space; the sphere + // test culls the corners. Same formula as the Triton / Metal kernels. + g.bhx = (int)std::ceil(r*c.fnx*c.inva); + g.bhy = (int)std::ceil(r*c.fny*c.invb); + g.bhz = (int)std::ceil(r*c.fnz*c.invc); + const scalar_t ax=xyz[3*a+0], ay=xyz[3*a+1], az=xyz[3*a+2]; + scalar_t fx=ax*im[0]+ay*im[1]+az*im[2]; + scalar_t fy=ax*im[3]+ay*im[4]+az*im[5]; + scalar_t fz=ax*im[6]+ay*im[7]+az*im[8]; + fx-=std::floor(fx); fy-=std::floor(fy); fz-=std::floor(fz); // wrap to [0,1) + g.cix=(int)std::nearbyint(fx*c.fnx); + g.ciy=(int)std::nearbyint(fy*c.fny); + g.ciz=(int)std::nearbyint(fz*c.fnz); + // Sub-voxel residual, taken to Cartesian: the sphere is centred HERE, on the + // atom, not on the anchor node. + const scalar_t sx=fx-(scalar_t)g.cix/c.fnx; + const scalar_t sy=fy-(scalar_t)g.ciy/c.fny; + const scalar_t sz=fz-(scalar_t)g.ciz/c.fnz; + g.w0x=fm[0]*sx+fm[1]*sy+fm[2]*sz; + g.w0y=fm[3]*sx+fm[4]*sy+fm[5]*sz; + g.w0z=fm[6]*sx+fm[7]*sy+fm[8]*sz; + return g; +} + +static inline int wrap_idx(int i, int n) { return ((i % n) + n) % n; } + +// Does any x-plane this atom touches fall in [xlo, xhi)? The plane set is the +// wrapped interval [cix-bhx, cix+bhx] mod nx, so at most two pieces. +static inline bool touches_x(int cix, int bhx, int nx, int64_t xlo, int64_t xhi) { + if (2*bhx+1 >= nx) return true; + const int lo = wrap_idx(cix-bhx, nx), hi = lo + 2*bhx; + if (hi < nx) return (lo < xhi) && (hi >= xlo); + return (lo < xhi) || ((hi - nx) >= xlo); +} + +// In-sphere oz run for one (ox,oy) column. r2(oz) is a convex quadratic in oz, so +// the in-sphere set is one contiguous run; solving for it keeps the innermost loop +// branchless. Widened by one voxel each way, with the exact r2 <= rc2 test retained +// as a 0/1 multiply, so the accepted voxel set is bit-identical to the GPU kernels' +// straight comparison over the full box. +template +static inline bool oz_run(const Cell& c, scalar_t qx, scalar_t qy, + scalar_t qz, scalar_t rc2, int bhz, int& zlo, int& zhi) { + const scalar_t Bq = (scalar_t)2*(qx*c.ucx+qy*c.ucy+qz*c.ucz); + const scalar_t Cq = qx*qx+qy*qy+qz*qz - rc2; + const scalar_t disc = Bq*Bq - (scalar_t)4*c.Auc*Cq; + if (disc < 0) return false; + const scalar_t sq = std::sqrt(disc), inv2A = (scalar_t)0.5/c.Auc; + zlo = (int)std::ceil((-Bq-sq)*inv2A) - 1; + zhi = (int)std::floor((-Bq+sq)*inv2A) + 1; + if (zlo < -bhz) zlo = -bhz; + if (zhi > bhz) zhi = bhz; + return zlo <= zhi; +} + +// =========================================================================== +// ISOTROPIC +// =========================================================================== +template +static void iso_fwd_impl(scalar_t* out, const scalar_t* xyz, const scalar_t* adp, + const scalar_t* occ, const scalar_t* Ac, const scalar_t* Bc, + const scalar_t* r2cut, const scalar_t* im, const scalar_t* fm, + int64_t n_at, int nx, int ny, int nz) +{ + const Cell c = make_cell(fm, im, nx, ny, nz); + // Output partitioned by x-plane: disjoint writes, so no atomics. + at::parallel_for(0, nx, 1, [&](int64_t xlo, int64_t xhi) { + for (int64_t a = 0; a < n_at; ++a) { + const Anchor g = make_anchor(c, xyz, fm, im, a, r2cut[a]); + if (!touches_x(g.cix, g.bhx, nx, xlo, xhi)) continue; + scalar_t Bt[5], An[5]; + const scalar_t bi=adp[a], oc=occ[a]; + for (int k=0;k<5;++k) { + scalar_t bt=(Bc[5*a+k]+bi)*(scalar_t)0.25; + bt = bt > (scalar_t)0.1 ? bt : (scalar_t)0.1; + Bt[k]=bt; + An[k]=Ac[5*a+k]*oc*K::PI_1P5/(bt*std::sqrt(bt)); + } + for (int ox=-g.bhx; ox<=g.bhx; ++ox) { + const int vix = wrap_idx(g.cix+ox, nx); + if (vix < xlo || vix >= xhi) continue; + const scalar_t fox=(scalar_t)ox; + const scalar_t px=fox*c.uax-g.w0x, py=fox*c.uay-g.w0y, pz=fox*c.uaz-g.w0z; + for (int oy=-g.bhy; oy<=g.bhy; ++oy) { + const scalar_t foy=(scalar_t)oy; + const scalar_t qx=px+foy*c.ubx, qy=py+foy*c.uby, qz=pz+foy*c.ubz; + int zlo, zhi; + if (!oz_run(c, qx, qy, qz, g.rc2, g.bhz, zlo, zhi)) continue; + scalar_t* row = out + ((int64_t)vix*ny + wrap_idx(g.ciy+oy, ny))*(int64_t)nz; + for (int oz=zlo; oz<=zhi; ++oz) { + const scalar_t foz=(scalar_t)oz; + const scalar_t wx=qx+foz*c.ucx, wy=qy+foz*c.ucy, wz=qz+foz*c.ucz; + const scalar_t r2=wx*wx+wy*wy+wz*wz; + const scalar_t keep = r2 <= g.rc2 ? (scalar_t)1 : (scalar_t)0; + scalar_t dens=0; + for (int k=0;k<5;++k) dens += An[k]*kexp(-K::PI_SQ*r2/Bt[k]); + row[wrap_idx(g.ciz+oz, nz)] += keep*dens; + } + } + } + } + }); +} + +template +static void iso_bwd_impl(scalar_t* g_xyz, scalar_t* g_adp, scalar_t* g_occ, + const scalar_t* go, const scalar_t* xyz, const scalar_t* adp, + const scalar_t* occ, const scalar_t* Ac, const scalar_t* Bc, + const scalar_t* r2cut, const scalar_t* im, const scalar_t* fm, + int64_t n_at, int nx, int ny, int nz) +{ + const Cell c = make_cell(fm, im, nx, ny, nz); + // Partitioned over atoms: each thread owns its atoms' gradient slots. + at::parallel_for(0, n_at, 1, [&](int64_t a0, int64_t a1) { + for (int64_t a = a0; a < a1; ++a) { + const Anchor g = make_anchor(c, xyz, fm, im, a, r2cut[a]); + scalar_t Bt[5], An[5], clampf[5]; + const scalar_t bi=adp[a], oc=occ[a]; + for (int k=0;k<5;++k) { + const scalar_t raw=(Bc[5*a+k]+bi)*(scalar_t)0.25; + const scalar_t bt = raw > (scalar_t)0.1 ? raw : (scalar_t)0.1; + Bt[k]=bt; + An[k]=Ac[5*a+k]*oc*K::PI_1P5/(bt*std::sqrt(bt)); + // in the clamp region d(Bt)/d(adp) = 0 + clampf[k] = raw > (scalar_t)0.1 ? (scalar_t)1 : (scalar_t)0; + } + scalar_t gx=0,gy=0,gz=0,gb=0,gocc=0; + for (int ox=-g.bhx; ox<=g.bhx; ++ox) { + const int vix = wrap_idx(g.cix+ox, nx); + const scalar_t fox=(scalar_t)ox; + const scalar_t px=fox*c.uax-g.w0x, py=fox*c.uay-g.w0y, pz=fox*c.uaz-g.w0z; + for (int oy=-g.bhy; oy<=g.bhy; ++oy) { + const scalar_t foy=(scalar_t)oy; + const scalar_t qx=px+foy*c.ubx, qy=py+foy*c.uby, qz=pz+foy*c.ubz; + int zlo, zhi; + if (!oz_run(c, qx, qy, qz, g.rc2, g.bhz, zlo, zhi)) continue; + const scalar_t* row = go + ((int64_t)vix*ny + wrap_idx(g.ciy+oy, ny))*(int64_t)nz; + for (int oz=zlo; oz<=zhi; ++oz) { + const scalar_t foz=(scalar_t)oz; + const scalar_t wx=qx+foz*c.ucx, wy=qy+foz*c.ucy, wz=qz+foz*c.ucz; + const scalar_t r2=wx*wx+wy*wy+wz*wz; + if (r2 > g.rc2) continue; + const scalar_t gout = row[wrap_idx(g.ciz+oz, nz)]; + scalar_t dens=0, coeff=0, dbs=0; + for (int k=0;k<5;++k) { + const scalar_t Ae = An[k]*kexp(-K::PI_SQ*r2/Bt[k]); + dens += Ae; + coeff += Ae/Bt[k]; + dbs += Ae*(-(scalar_t)1.5/Bt[k] + + K::PI_SQ*r2/(Bt[k]*Bt[k]))*clampf[k]; + } + const scalar_t s = gout*(scalar_t)2*K::PI_SQ*coeff; + gx += s*wx; gy += s*wy; gz += s*wz; + gb += gout*(scalar_t)0.25*dbs; + gocc += gout*dens; + } + } + } + g_xyz[3*a+0]=gx; g_xyz[3*a+1]=gy; g_xyz[3*a+2]=gz; + g_adp[a]=gb; + g_occ[a]= oc != (scalar_t)0 ? gocc/oc : (scalar_t)0; + } + }); +} + +// =========================================================================== +// ANISOTROPIC. M_g = (B_g*I + 8*pi^2*U)/4, inverted analytically; density uses +// the Mahalanobis form q = w^T Minv w, the cutoff stays the Euclidean sphere. +// =========================================================================== +template +static inline void aniso_minv(const scalar_t* Bc, const scalar_t* u, int64_t a, + const scalar_t* Ac, scalar_t oc, + scalar_t* p00, scalar_t* p11, scalar_t* p22, + scalar_t* p01, scalar_t* p02, scalar_t* p12, + scalar_t* An) +{ + const scalar_t T = K::TWO_PI_SQ; // 8*pi^2 / 4 + const scalar_t u11=u[6*a+0],u22=u[6*a+1],u33=u[6*a+2], + u12=u[6*a+3],u13=u[6*a+4],u23=u[6*a+5]; + const scalar_t md=T*u12, me=T*u13, mf=T*u23; + for (int k=0;k<5;++k) { + const scalar_t Bg=Bc[5*a+k]; + const scalar_t ma=(scalar_t)0.25*Bg+T*u11; + const scalar_t mb=(scalar_t)0.25*Bg+T*u22; + const scalar_t mc=(scalar_t)0.25*Bg+T*u33; + const scalar_t det=ma*(mb*mc-mf*mf)-md*(md*mc-me*mf)+me*(md*mf-me*mb); + const scalar_t inv=(scalar_t)1/det; + p00[k]=(mb*mc-mf*mf)*inv; p11[k]=(ma*mc-me*me)*inv; p22[k]=(ma*mb-md*md)*inv; + p01[k]=(me*mf-md*mc)*inv; p02[k]=(md*mf-me*mb)*inv; p12[k]=(md*me-ma*mf)*inv; + const scalar_t d = det > (scalar_t)1e-10 ? det : (scalar_t)1e-10; + An[k]=Ac[5*a+k]*oc*K::PI_1P5/std::sqrt(d); + } +} + +template +static void aniso_fwd_impl(scalar_t* out, const scalar_t* xyz, const scalar_t* u, + const scalar_t* occ, const scalar_t* Ac, const scalar_t* Bc, + const scalar_t* r2cut, const scalar_t* im, const scalar_t* fm, + int64_t n_at, int nx, int ny, int nz) +{ + const Cell c = make_cell(fm, im, nx, ny, nz); + at::parallel_for(0, nx, 1, [&](int64_t xlo, int64_t xhi) { + for (int64_t a = 0; a < n_at; ++a) { + const Anchor g = make_anchor(c, xyz, fm, im, a, r2cut[a]); + if (!touches_x(g.cix, g.bhx, nx, xlo, xhi)) continue; + scalar_t p00[5],p11[5],p22[5],p01[5],p02[5],p12[5],An[5]; + aniso_minv(Bc, u, a, Ac, occ[a], p00,p11,p22,p01,p02,p12,An); + for (int ox=-g.bhx; ox<=g.bhx; ++ox) { + const int vix = wrap_idx(g.cix+ox, nx); + if (vix < xlo || vix >= xhi) continue; + const scalar_t fox=(scalar_t)ox; + const scalar_t px=fox*c.uax-g.w0x, py=fox*c.uay-g.w0y, pz=fox*c.uaz-g.w0z; + for (int oy=-g.bhy; oy<=g.bhy; ++oy) { + const scalar_t foy=(scalar_t)oy; + const scalar_t qx=px+foy*c.ubx, qy=py+foy*c.uby, qz=pz+foy*c.ubz; + int zlo, zhi; + if (!oz_run(c, qx, qy, qz, g.rc2, g.bhz, zlo, zhi)) continue; + scalar_t* row = out + ((int64_t)vix*ny + wrap_idx(g.ciy+oy, ny))*(int64_t)nz; + for (int oz=zlo; oz<=zhi; ++oz) { + const scalar_t foz=(scalar_t)oz; + const scalar_t wx=qx+foz*c.ucx, wy=qy+foz*c.ucy, wz=qz+foz*c.ucz; + if (wx*wx+wy*wy+wz*wz > g.rc2) continue; + const scalar_t xx=wx*wx, yy=wy*wy, zz=wz*wz; + const scalar_t xy=wx*wy, xz=wx*wz, yz=wy*wz; + scalar_t dens=0; + for (int k=0;k<5;++k) { + const scalar_t q = p00[k]*xx+p11[k]*yy+p22[k]*zz + + (scalar_t)2*(p01[k]*xy+p02[k]*xz+p12[k]*yz); + dens += An[k]*kexp(-K::PI_SQ*q); + } + row[wrap_idx(g.ciz+oz, nz)] += dens; + } + } + } + } + }); +} + +template +static void aniso_bwd_impl(scalar_t* g_xyz, scalar_t* g_u, scalar_t* g_occ, + const scalar_t* go, const scalar_t* xyz, const scalar_t* u, + const scalar_t* occ, const scalar_t* Ac, const scalar_t* Bc, + const scalar_t* r2cut, const scalar_t* im, const scalar_t* fm, + int64_t n_at, int nx, int ny, int nz) +{ + const Cell c = make_cell(fm, im, nx, ny, nz); + at::parallel_for(0, n_at, 1, [&](int64_t a0, int64_t a1) { + for (int64_t a = a0; a < a1; ++a) { + const Anchor g = make_anchor(c, xyz, fm, im, a, r2cut[a]); + const scalar_t oc = occ[a]; + scalar_t p00[5],p11[5],p22[5],p01[5],p02[5],p12[5],An[5]; + aniso_minv(Bc, u, a, Ac, oc, p00,p11,p22,p01,p02,p12,An); + scalar_t gx=0,gy=0,gz=0,gu0=0,gu1=0,gu2=0,gu3=0,gu4=0,gu5=0,gocc=0; + for (int ox=-g.bhx; ox<=g.bhx; ++ox) { + const int vix = wrap_idx(g.cix+ox, nx); + const scalar_t fox=(scalar_t)ox; + const scalar_t px=fox*c.uax-g.w0x, py=fox*c.uay-g.w0y, pz=fox*c.uaz-g.w0z; + for (int oy=-g.bhy; oy<=g.bhy; ++oy) { + const scalar_t foy=(scalar_t)oy; + const scalar_t qx=px+foy*c.ubx, qy=py+foy*c.uby, qz=pz+foy*c.ubz; + int zlo, zhi; + if (!oz_run(c, qx, qy, qz, g.rc2, g.bhz, zlo, zhi)) continue; + const scalar_t* row = go + ((int64_t)vix*ny + wrap_idx(g.ciy+oy, ny))*(int64_t)nz; + for (int oz=zlo; oz<=zhi; ++oz) { + const scalar_t foz=(scalar_t)oz; + const scalar_t wx=qx+foz*c.ucx, wy=qy+foz*c.ucy, wz=qz+foz*c.ucz; + if (wx*wx+wy*wy+wz*wz > g.rc2) continue; + const scalar_t gout = row[wrap_idx(g.ciz+oz, nz)]; + scalar_t dens=0,sx=0,sy=0,sz=0,s0=0,s1=0,s2=0,s3=0,s4=0,s5=0; + for (int k=0;k<5;++k) { + const scalar_t vx=p00[k]*wx+p01[k]*wy+p02[k]*wz; + const scalar_t vy=p01[k]*wx+p11[k]*wy+p12[k]*wz; + const scalar_t vz=p02[k]*wx+p12[k]*wy+p22[k]*wz; + const scalar_t q = wx*vx+wy*vy+wz*vz; + const scalar_t dg = An[k]*kexp(-K::PI_SQ*q); + dens+=dg; sx+=dg*vx; sy+=dg*vy; sz+=dg*vz; + const scalar_t P = K::PI_SQ; + s0+=dg*(-(scalar_t)0.5*p00[k]+P*vx*vx); + s1+=dg*(-(scalar_t)0.5*p11[k]+P*vy*vy); + s2+=dg*(-(scalar_t)0.5*p22[k]+P*vz*vz); + s3+=dg*(-(scalar_t)0.5*p01[k]+P*vx*vy); + s4+=dg*(-(scalar_t)0.5*p02[k]+P*vx*vz); + s5+=dg*(-(scalar_t)0.5*p12[k]+P*vy*vz); + } + const scalar_t s2pi = gout*(scalar_t)2*K::PI_SQ; + const scalar_t s4pi = gout*(scalar_t)4*K::PI_SQ; + gx+=s2pi*sx; gy+=s2pi*sy; gz+=s2pi*sz; + gu0+=s2pi*s0; gu1+=s2pi*s1; gu2+=s2pi*s2; // diagonal U + gu3+=s4pi*s3; gu4+=s4pi*s4; gu5+=s4pi*s5; // off-diagonal U + gocc+=gout*dens; + } + } + } + g_xyz[3*a+0]=gx; g_xyz[3*a+1]=gy; g_xyz[3*a+2]=gz; + g_u[6*a+0]=gu0; g_u[6*a+1]=gu1; g_u[6*a+2]=gu2; + g_u[6*a+3]=gu3; g_u[6*a+4]=gu4; g_u[6*a+5]=gu5; + g_occ[a]= oc != (scalar_t)0 ? gocc/oc : (scalar_t)0; + } + }); +} + +// =========================================================================== +// Bindings. Tensors arrive contiguous and pre-validated from Python. +// =========================================================================== +#define PTR(T, t) (t).data_ptr() + +void iso_fwd(torch::Tensor out, torch::Tensor xyz, torch::Tensor adp, + torch::Tensor occ, torch::Tensor A, torch::Tensor B, + torch::Tensor r2cut, torch::Tensor inv_frac, torch::Tensor frac, + int64_t nx, int64_t ny, int64_t nz) +{ + AT_DISPATCH_FLOATING_TYPES(out.scalar_type(), "iso_fwd", [&] { + iso_fwd_impl(PTR(scalar_t,out), PTR(scalar_t,xyz), PTR(scalar_t,adp), + PTR(scalar_t,occ), PTR(scalar_t,A), PTR(scalar_t,B), PTR(scalar_t,r2cut), + PTR(scalar_t,inv_frac), PTR(scalar_t,frac), xyz.size(0), nx, ny, nz); + }); +} + +void iso_bwd(torch::Tensor g_xyz, torch::Tensor g_adp, torch::Tensor g_occ, + torch::Tensor go, torch::Tensor xyz, torch::Tensor adp, + torch::Tensor occ, torch::Tensor A, torch::Tensor B, + torch::Tensor r2cut, torch::Tensor inv_frac, torch::Tensor frac, + int64_t nx, int64_t ny, int64_t nz) +{ + AT_DISPATCH_FLOATING_TYPES(go.scalar_type(), "iso_bwd", [&] { + iso_bwd_impl(PTR(scalar_t,g_xyz), PTR(scalar_t,g_adp), + PTR(scalar_t,g_occ), PTR(scalar_t,go), PTR(scalar_t,xyz), + PTR(scalar_t,adp), PTR(scalar_t,occ), PTR(scalar_t,A), PTR(scalar_t,B), + PTR(scalar_t,r2cut), PTR(scalar_t,inv_frac), PTR(scalar_t,frac), + xyz.size(0), nx, ny, nz); + }); +} + +void aniso_fwd(torch::Tensor out, torch::Tensor xyz, torch::Tensor u, + torch::Tensor occ, torch::Tensor A, torch::Tensor B, + torch::Tensor r2cut, torch::Tensor inv_frac, torch::Tensor frac, + int64_t nx, int64_t ny, int64_t nz) +{ + AT_DISPATCH_FLOATING_TYPES(out.scalar_type(), "aniso_fwd", [&] { + aniso_fwd_impl(PTR(scalar_t,out), PTR(scalar_t,xyz), PTR(scalar_t,u), + PTR(scalar_t,occ), PTR(scalar_t,A), PTR(scalar_t,B), PTR(scalar_t,r2cut), + PTR(scalar_t,inv_frac), PTR(scalar_t,frac), xyz.size(0), nx, ny, nz); + }); +} + +void aniso_bwd(torch::Tensor g_xyz, torch::Tensor g_u, torch::Tensor g_occ, + torch::Tensor go, torch::Tensor xyz, torch::Tensor u, + torch::Tensor occ, torch::Tensor A, torch::Tensor B, + torch::Tensor r2cut, torch::Tensor inv_frac, torch::Tensor frac, + int64_t nx, int64_t ny, int64_t nz) +{ + AT_DISPATCH_FLOATING_TYPES(go.scalar_type(), "aniso_bwd", [&] { + aniso_bwd_impl(PTR(scalar_t,g_xyz), PTR(scalar_t,g_u), + PTR(scalar_t,g_occ), PTR(scalar_t,go), PTR(scalar_t,xyz), + PTR(scalar_t,u), PTR(scalar_t,occ), PTR(scalar_t,A), PTR(scalar_t,B), + PTR(scalar_t,r2cut), PTR(scalar_t,inv_frac), PTR(scalar_t,frac), + xyz.size(0), nx, ny, nz); + }); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("iso_fwd", &iso_fwd, "Fused isotropic spherical-cutoff splat (forward)"); + m.def("iso_bwd", &iso_bwd, "Fused isotropic spherical-cutoff splat (backward)"); + m.def("aniso_fwd", &aniso_fwd, "Fused anisotropic spherical-cutoff splat (forward)"); + m.def("aniso_bwd", &aniso_bwd, "Fused anisotropic spherical-cutoff splat (backward)"); +} +""" + +# --------------------------------------------------------------------------- +# Lazy compilation, attempted at most once per process. Mirrors +# ``scatter.py::_get_module``: any failure returns None so the caller degrades to +# the portable plain splat rather than dying. +# --------------------------------------------------------------------------- +_module = None +_module_failed = False +_module_error: Optional[Tuple[str, str]] = None + + +def _get_module(): + """The compiled fused-splat extension, or None if it could not be built.""" + global _module, _module_failed, _module_error + if _module is not None: + return _module + if _module_failed: + return None + _module, _module_error = build_extension("sphere_splat", _CPP_SRC) + if _module is None: + _module_failed = True + return _module + + +def why_unavailable() -> Optional[str]: + """``None`` if the fused CPU splat is usable, else why it is not. + + The single availability probe for this backend -- the shape every backend implements, + consumed by :mod:`torchref.utils.backends`. A missing compiler and a compile error are + different problems, and the captured diagnostic is the only thing that separates them. + """ + if _get_module() is not None: + return None + reason = _module_error[0] if _module_error else "unknown reason" + return ( + f"the fused CPU sphere_splat extension is not available ({reason}); see " + "torchref.base.electron_density.kernels.cpu.sphere_splat.last_error()" + ) + + +def sphere_splat_available() -> bool: + """Whether the fused CPU splat compiled and is ready to dispatch. + + Derived from :func:`why_unavailable` rather than re-testing, so there is one + availability check here, not two that can drift. + """ + return why_unavailable() is None + + + + +def warmup() -> bool: + """Eagerly compile, to move the one-time cost off the first refinement step.""" + return _get_module() is not None + + +def clear_cache() -> None: + """Forget the compiled module and failure state (rebuilt on next use).""" + global _module, _module_failed, _module_error + _module = None + _module_failed = False + _module_error = None + + +def last_error() -> Optional[Tuple[str, str]]: + """The ``(message, traceback)`` of the last build failure, if any.""" + return _module_error + + +def _require_module(): + mod = _get_module() + if mod is None: + err = _module_error[0] if _module_error else "unknown reason" + raise RuntimeError( + f"fused CPU sphere_splat extension not available ({err}). See " + "torchref.base.electron_density.kernels.cpu.sphere_splat.last_error()." + ) + return mod + + +def _double_backward_vjp(plain_fn, ctx, grad_out, leaves, statics, r2cut): + """Recompute this VJP through the portable differentiable splat. + + The C++ backward is a closed-form first-order formula with no autograd graph, so + on its own it cannot supply a second derivative -- the same limitation the CUDA + and Metal kernels have. Rather than lose double backward on the CPU default path + (``test_kernel_fixes.py`` covers it, and a *silently wrong* Hessian was a real bug + here once), detect the double-backward context and re-derive the identical VJP + from the portable splat, which is built from differentiable ops. + + ``torch.is_grad_enabled()`` is the detector: autograd runs ``backward`` under + ``no_grad`` unless the caller passed ``create_graph=True``. ``grad_out``'s own + ``requires_grad`` is *not* a usable signal -- at the top of a + ``create_graph=True`` backward it is a plain ``ones`` tensor. + + The gradients are taken w.r.t. the **saved** leaves, not detached copies, so the + returned VJP stays connected to the caller's graph -- detaching here would + silently drop the second-order term, which is precisely the failure mode this + guards against. Both paths implement the same truncation contract, so the + first-order values agree to float noise; the only cost is that a Hessian + workflow runs at the portable splat's speed. + """ + A, B, inv_frac, frac = statics + # Only the leaves that actually require grad may be differentiated; asking for + # the others raises regardless of allow_unused. + wanted = [t for t in leaves if t.requires_grad] + out = [None] * len(leaves) + if wanted: + zeros = torch.zeros( + ctx.grid_shape, dtype=grad_out.dtype, device=grad_out.device + ) + with torch.enable_grad(): + dm = plain_fn(zeros, *leaves, A, B, inv_frac, frac, r2cut.sqrt()) + grads = torch.autograd.grad( + dm, wanted, grad_out, create_graph=True, allow_unused=True + ) + it = iter(grads) + out = [next(it) if t.requires_grad else None for t in leaves] + # forward returned density_map + splat -> grad wrt density_map is the identity + return (grad_out, *out) + (None,) * 5 + + +def _prep(density_map, xyz, radius_per_atom, *tensors): + """Shared validation + contiguity for both entry points.""" + if density_map.device.type != "cpu": + raise ValueError( + f"sphere_splat is a CPU kernel; got device {density_map.device}" + ) + dtype = density_map.dtype + if dtype not in (torch.float32, torch.float64): + raise ValueError(f"sphere_splat supports float32/float64, got {dtype}") + for t in (xyz, radius_per_atom) + tensors: + if t.dtype != dtype: + raise ValueError( + f"every input must match density_map.dtype ({dtype}); got {t.dtype}" + ) + r2cut = (radius_per_atom * radius_per_atom).detach().contiguous() + return dtype, r2cut + + +class _FusedIsoSplat(torch.autograd.Function): + """Isotropic fused CPU splat: returns ``density_map + splat``.""" + + @staticmethod + def forward(ctx, density_map, xyz, adp, occ, A, B, r2cut, inv_frac, frac): + mod = _require_module() + nx, ny, nz = (int(s) for s in density_map.shape) + out = density_map.contiguous().clone() + if xyz.shape[0] > 0: + mod.iso_fwd( + out.view(-1), + xyz.detach().contiguous(), adp.detach().contiguous(), + occ.detach().contiguous(), A.contiguous(), B.contiguous(), + r2cut, inv_frac.contiguous().view(-1), frac.contiguous().view(-1), + nx, ny, nz, + ) + ctx.save_for_backward(xyz, adp, occ, A, B, r2cut, inv_frac, frac) + ctx.grid_shape = (nx, ny, nz) + return out + + @staticmethod + def backward(ctx, grad_out): + xyz, adp, occ, A, B, r2cut, inv_frac, frac = ctx.saved_tensors + nx, ny, nz = ctx.grid_shape + if torch.is_grad_enabled(): # create_graph=True -> need a differentiable VJP + from torchref.base.electron_density.kernels.cpu.variable_radius import ( + add_isotropic_plain_var, + ) + + return _double_backward_vjp( + add_isotropic_plain_var, ctx, grad_out, + (xyz, adp, occ), (A, B, inv_frac, frac), r2cut, + ) + g_xyz = torch.zeros_like(xyz) + g_adp = torch.zeros_like(adp) + g_occ = torch.zeros_like(occ) + if xyz.shape[0] > 0: + _require_module().iso_bwd( + g_xyz.view(-1), g_adp, g_occ, + grad_out.contiguous().view(-1), + xyz.detach().contiguous(), adp.detach().contiguous(), + occ.detach().contiguous(), A.contiguous(), B.contiguous(), + r2cut, inv_frac.contiguous().view(-1), frac.contiguous().view(-1), + nx, ny, nz, + ) + # out = density_map + splat, so grad wrt density_map is the identity. + return (grad_out, g_xyz, g_adp, g_occ, None, None, None, None, None) + + +class _FusedAnisoSplat(torch.autograd.Function): + """Anisotropic fused CPU splat: returns ``density_map + splat``.""" + + @staticmethod + def forward(ctx, density_map, xyz, u, occ, A, B, r2cut, inv_frac, frac): + mod = _require_module() + nx, ny, nz = (int(s) for s in density_map.shape) + out = density_map.contiguous().clone() + if xyz.shape[0] > 0: + mod.aniso_fwd( + out.view(-1), + xyz.detach().contiguous(), u.detach().contiguous(), + occ.detach().contiguous(), A.contiguous(), B.contiguous(), + r2cut, inv_frac.contiguous().view(-1), frac.contiguous().view(-1), + nx, ny, nz, + ) + ctx.save_for_backward(xyz, u, occ, A, B, r2cut, inv_frac, frac) + ctx.grid_shape = (nx, ny, nz) + return out + + @staticmethod + def backward(ctx, grad_out): + xyz, u, occ, A, B, r2cut, inv_frac, frac = ctx.saved_tensors + nx, ny, nz = ctx.grid_shape + if torch.is_grad_enabled(): # see _FusedIsoSplat.backward + from torchref.base.electron_density.kernels.cpu.variable_radius import ( + add_anisotropic_plain_var, + ) + + return _double_backward_vjp( + add_anisotropic_plain_var, ctx, grad_out, + (xyz, u, occ), (A, B, inv_frac, frac), r2cut, + ) + g_xyz = torch.zeros_like(xyz) + g_u = torch.zeros_like(u) + g_occ = torch.zeros_like(occ) + if xyz.shape[0] > 0: + _require_module().aniso_bwd( + g_xyz.view(-1), g_u.view(-1), g_occ, + grad_out.contiguous().view(-1), + xyz.detach().contiguous(), u.detach().contiguous(), + occ.detach().contiguous(), A.contiguous(), B.contiguous(), + r2cut, inv_frac.contiguous().view(-1), frac.contiguous().view(-1), + nx, ny, nz, + ) + return (grad_out, g_xyz, g_u, g_occ, None, None, None, None, None) + + +def add_isotropic_cpu_sphere_var( + density_map, xyz, adp, occ, A, B, inv_frac_matrix, frac_matrix, radius_per_atom +): + """Fused isotropic spherical-cutoff splat; adds into ``density_map``. + + Parameters + ---------- + density_map : torch.Tensor + Running density map, shape (nx, ny, nz), CPU float32/float64. Not mutated; + the sum is returned. + xyz : torch.Tensor + Cartesian atom positions, shape (n, 3). + adp : torch.Tensor + Isotropic B-factors, shape (n,). + occ : torch.Tensor + Occupancies, shape (n,). + A, B : torch.Tensor + ITC92 amplitudes / widths, shape (n, 5). + inv_frac_matrix, frac_matrix : torch.Tensor + Cartesian<->fractional matrices, shape (3, 3). The truncation box is + derived from these, so no ``voxel_size`` argument is needed. + radius_per_atom : torch.Tensor + Per-atom cutoff radius in Angstrom, shape (n,), from + :func:`radius_policy.per_atom_radius_iso`. Used raw -- no grid-dependent + requantization, so the cutoff means the same thing at any grid sampling. + + Returns + ------- + torch.Tensor + ``density_map + splat``, shape (nx, ny, nz). + """ + _, r2cut = _prep(density_map, xyz, radius_per_atom, adp, occ, A, B) + return _FusedIsoSplat.apply( + density_map, xyz, adp, occ, A, B, r2cut, inv_frac_matrix, frac_matrix + ) + + +def add_anisotropic_cpu_sphere_var( + density_map, xyz, u, occ, A, B, inv_frac_matrix, frac_matrix, radius_per_atom +): + """Fused anisotropic spherical-cutoff splat; adds into ``density_map``. + + Identical contract to :func:`add_isotropic_cpu_sphere_var`, but ``u`` carries + the 6 anisotropic components ``[U11, U22, U33, U12, U13, U23]`` and the density + is the full 3D Gaussian ``exp(-pi^2 w^T Minv w)`` with + ``M_g = (B_g*I + 8*pi^2*U)/4``. The *cutoff* remains the Euclidean sphere at + ``radius_per_atom`` (the ellipsoid's isotropic bounding radius, from + :func:`radius_policy.per_atom_radius_aniso`) -- matching the CUDA and Metal + kernels, which likewise cull on Euclidean distance and evaluate the + Mahalanobis form. + """ + _, r2cut = _prep(density_map, xyz, radius_per_atom, u, occ, A, B) + return _FusedAnisoSplat.apply( + density_map, xyz, u, occ, A, B, r2cut, inv_frac_matrix, frac_matrix + ) diff --git a/torchref/base/electron_density/kernels/cpu/variable_radius.py b/torchref/base/electron_density/kernels/cpu/variable_radius.py index 0be1a8da..595dea1d 100644 --- a/torchref/base/electron_density/kernels/cpu/variable_radius.py +++ b/torchref/base/electron_density/kernels/cpu/variable_radius.py @@ -1,40 +1,42 @@ -"""CPU per-atom variable-radius density splatting (grouped-separable / fused / aniso). - -Each atom is truncated at its own ``N_sigma * sigma_eff`` radius instead of a -single structure-wide radius. The separable CPU path factorizes -``exp(-alpha * r^T G r)`` into 1D per-axis Gaussians (O(r) exp calls, not O(r^3)) -which needs a *uniform* box per batch, so a per-atom radius forces grouping atoms -by box. Two choices keep the grouping cheap: - -* **Bucket by integer box size** ``box_radius = ceil(r_i / min_voxel)`` (NOT the - nominal 0.25-A radius): the kernels truncate by the box, so atoms rounding to - the same box are identical and share one launch. -* **One global sort by (box, center_1d)** -> each bucket is a contiguous, - cache-sorted slice; then **chunk atoms within a bucket** for L3 locality. - -Work drops from ``N * max_box^3`` to ``sum_bucket n * box^3`` while the -factorization is preserved. The per-voxel cores (``_separable_density``, -``_aniso_density_cube``) and the structured scatter are reused verbatim from the -single-radius kernels, so a single-radius plan reproduces them bit-for-bit. No -``torch.compile``. - -These functions ADD into the supplied ``density_map`` (so the isotropic and -anisotropic passes accumulate into the same map) and are autograd-connected in -xyz / adp / u / occ. +"""Portable per-atom variable-radius density splatting. + +Reached by ``force_portable`` on any device, by CUDA/MPS float64, and whenever the fused +C++ kernel could not be built. Plain ``scatter_add`` only, so it runs on every device, +supports float64, and is double-differentiable -- which makes it the reference the +accelerator kernels are checked against. + +One truncation contract, shared with the Triton, Metal and fused-CPU kernels, so AUTO +and EAGER agree to float noise on every device: + + voxel v gets atom i's density iff ``||w||^2 <= r_i^2``, where ``w`` is the Cartesian + atom->voxel vector (sphere centred on the ATOM, not on its anchor node) and ``r_i`` + is the raw radius_policy radius, + +enumerated over the triclinic-correct per-axis box +``ceil(r * n_axis * ||inv_frac row_axis||)``. See ``sphere_splat.py`` for the canonical +statement. + +These used to diverge from that in three ways at once: the iso path selected voxels by +``||offset * voxel_size||`` -- a diagonal metric, wrong for any non-orthogonal cell -- +measured from the anchor node rather than from the atom, at a radius rounded up to a +whole voxel; and the aniso path splatted a full cube. On a beta=115 deg cell that +mis-selected ~12% of each sphere's voxels, a 5e-3 rel L2 map error, i.e. larger than the +1.7e-3 truncation error the cutoff exists to deliver. + +Out-of-sphere voxels are zeroed rather than dropped, keeping the box dense so one +``scatter_add`` covers the chunk. That wastes some writes, which is the right trade for a +portable reference. + +These functions ADD into the supplied ``density_map`` (so the isotropic and anisotropic +passes accumulate into the same map) and are autograd-connected in xyz / adp / u / occ. """ from __future__ import annotations import math -from typing import List, Tuple import torch -from torchref.config import dtypes -from torchref.base.electron_density.kernels.offsets import _get_radius_offsets -from torchref.base.electron_density.kernels.cpu.separable import _separable_density -from torchref.base.electron_density.kernels.cpu.aniso import _aniso_density_cube -from torchref.base.electron_density.kernels.cpu.scatter_dispatch import _do_structured_scatter from torchref.base.electron_density.radius_policy import _u6_to_u3 _PI = math.pi @@ -44,277 +46,154 @@ _CHUNK = 1024 -# ========================================================================= -# Shared grouping helpers -# ========================================================================= -def _cross_term_flags(G: torch.Tensor) -> Tuple[bool, bool, bool]: - tol = 1e-3 * torch.norm(torch.diagonal(G)) - return ( - bool(torch.abs(G[0, 1]) > tol), - bool(torch.abs(G[0, 2]) > tol), - bool(torch.abs(G[1, 2]) > tol), - ) - - -def _box_radius_per_atom(radius: torch.Tensor, voxel_size: torch.Tensor) -> torch.Tensor: - """Integer box radius per atom = ceil(r_i / min_voxel) -- the true bucket key.""" - min_voxel = float(voxel_size.min()) - return torch.ceil(radius / min_voxel).to(torch.int64) - +def _bucket_by_radius(radius: torch.Tensor, center_1d: torch.Tensor): + """Sort atoms by (radius, center_1d); return ``(order, [(radius, start, end)])``. -def _bucket_by_box(box_radius: torch.Tensor, center_1d: torch.Tensor): - """Sort atoms by (box_radius, center_1d); return (order, [(box_radius, start, end)]). - - ONE global sort -> each distinct-box bucket is a contiguous, cache-sorted slice. + Buckets on the radius itself rather than a voxel-derived box size: the policy + quantizes to 0.25 A and clamps to [2, 7], so there are at most 21 distinct + values and the per-axis box follows from the radius alone. Sorting by 1D centre + within a bucket keeps the scatter cache-friendly. """ - uniq = torch.unique(box_radius) order_parts, spans, cursor = [], [], 0 - for b in uniq.tolist(): - idx = (box_radius == b).nonzero(as_tuple=True)[0] + for r in torch.unique(radius).tolist(): + idx = (radius == r).nonzero(as_tuple=True)[0] idx = idx[torch.argsort(center_1d[idx])] order_parts.append(idx) - spans.append((int(b), cursor, cursor + idx.numel())) + spans.append((float(r), cursor, cursor + idx.numel())) cursor += idx.numel() order = (torch.cat(order_parts) if order_parts - else box_radius.new_zeros(0, dtype=torch.long)) + else torch.zeros(0, dtype=torch.long, device=radius.device)) return order, spans -def _axis_offsets(box_radius: int, device, int_dtype): - return torch.arange(-box_radius, box_radius + 1, device=device).to(int_dtype) - +def _axis_half_widths(r: float, inv_frac: torch.Tensor, grid_dims): + """``ceil(r * n_axis * ||inv_frac row_axis||)`` -- the kernels' enumeration box. -# ========================================================================= -# Isotropic separable (Engine.AUTO CPU path) -# ========================================================================= -def _splat_chunked(density_flat, map_size, G, flags, inv_grid, grid_dims, device, dtype, - xyz_frac, center_idx, alpha, A_norm, spans, chunk=_CHUNK): - """Bucket loop (per box size) x chunk loop: factorized cube + structured scatter.""" - nx, ny, nz = grid_dims - ny_nz = ny * nz - for box_radius, b0, b1 in spans: - axis = _axis_offsets(box_radius, device, dtypes.int) - axis_frac = axis.to(dtype).unsqueeze(0) * inv_grid.unsqueeze(1) - for s in range(b0, b1, chunk): - e = min(s + chunk, b1) - ci = center_idx[s:e] - center_frac = ci.to(dtype) * inv_grid - sub = xyz_frac[s:e] - center_frac - d_frac = axis_frac.unsqueeze(0) - sub.unsqueeze(2) - d_frac = d_frac - torch.round(d_frac) - cube = _separable_density(d_frac, alpha[s:e], A_norm[s:e], G, *flags) - wa = (ci[:, 0:1] + axis.unsqueeze(0)) % nx * ny_nz - wb = (ci[:, 1:2] + axis.unsqueeze(0)) % ny * nz - wc = (ci[:, 2:3] + axis.unsqueeze(0)) % nz - wbwc = wb.unsqueeze(2) + wc.unsqueeze(1) - density_flat = _do_structured_scatter(cube, wa, wbwc, density_flat, map_size) - return density_flat - - -def _iso_setup(xyz, adp, occ, A, B, inv_frac, frac, grid_shape, voxel_size, - radius_per_atom): - """Per-atom + shared setup; returns sorted tensors + bucket spans.""" - device, dtype = xyz.device, xyz.dtype - nx, ny, nz = grid_shape - ny_nz = ny * nz - grid_f = torch.tensor(grid_shape, device=device, dtype=dtype) - inv_grid = 1.0 / grid_f - - G = frac.T @ frac - flags = _cross_term_flags(G) - - xyz_frac = xyz @ inv_frac.T - center_idx = torch.round((xyz_frac % 1.0) * grid_f).to(dtypes.int) - B_total = ((B + adp[:, None]) * 0.25).clamp(min=0.1) - A_norm = A * occ[:, None] * _PI_1P5 / (B_total * torch.sqrt(B_total)) - alpha = _PI_SQ / B_total - - box_radius = _box_radius_per_atom(radius_per_atom, voxel_size) - center_1d = (center_idx[:, 0].long() * ny_nz - + center_idx[:, 1].long() * nz + center_idx[:, 2].long()) - order, spans = _bucket_by_box(box_radius, center_1d) - return dict( - G=G, flags=flags, inv_grid=inv_grid, grid_dims=(nx, ny, nz), - xyz_frac=xyz_frac[order], center_idx=center_idx[order], - alpha=alpha[order], A_norm=A_norm[order], B_total=B_total[order], - spans=spans, + The norm is taken in float64 **on the CPU**, never on the input's device: this path + also serves ``force_portable`` on MPS, which has no float64, and ``.double()`` in place + raises there. Hopping a 3x3 matrix to the CPU is free, and float64 matters because the + result feeds a ``ceil`` -- a value landing a hair under an integer in float32 would + shrink the box by one voxel and silently clip the sphere. + """ + row_norms = torch.linalg.norm(inv_frac.detach().cpu().double(), dim=1) + return tuple( + int(math.ceil(r * float(n) * float(row_norms[i]))) + for i, n in enumerate(grid_dims) ) -def add_isotropic_cpu_separable_var(density_map, xyz, adp, occ, A, B, - inv_frac_matrix, frac_matrix, grid_shape_tuple, - voxel_size, radius_per_atom): - """Variable-radius grouped-separable isotropic splat; adds into ``density_map``.""" - nx, ny, nz = grid_shape_tuple - map_size = nx * ny * nz - st = _iso_setup(xyz, adp, occ, A, B, inv_frac_matrix, frac_matrix, - grid_shape_tuple, voxel_size, radius_per_atom) - density_flat = density_map.reshape(-1) - density_flat = _splat_chunked( - density_flat, map_size, st["G"], st["flags"], st["inv_grid"], - st["grid_dims"], xyz.device, xyz.dtype, - st["xyz_frac"], st["center_idx"], st["alpha"], st["A_norm"], st["spans"], - ) - return density_flat.view(nx, ny, nz) +def _box_offsets(bh, frac, grid_dims, device, dtype): + """Per-axis box offsets plus their Cartesian displacements. + + Returns ``(offsets (R,3) int64, off_cart (R,3))`` with + ``off_cart = frac @ (offset / n)`` -- the same ``ox*u_a + oy*u_b + oz*u_c`` the + kernels accumulate. + """ + axes = [torch.arange(-b, b + 1, device=device) for b in bh] + gx, gy, gz = torch.meshgrid(*axes, indexing="ij") + offsets = torch.stack((gx, gy, gz), dim=-1).reshape(-1, 3) + inv_grid = 1.0 / torch.tensor(grid_dims, device=device, dtype=dtype) + return offsets, (offsets.to(dtype) * inv_grid) @ frac.T -# ========================================================================= -# Isotropic fused sphere (Engine.EAGER CPU path) -# ========================================================================= -def _splat_chunked_sphere(density_flat, G, inv_grid, grid_dims, voxel_size, device, dtype, - xyz_frac, center_idx, B_total, A_norm, spans, chunk=_CHUNK): - """Per box-bucket x chunk: sphere offsets + r^2 + 5-Gaussian + plain scatter_add_.""" - grid_shape = torch.tensor(grid_dims, device=device) - min_voxel = float(voxel_size.min()) +def _canonical_setup(xyz, inv_frac, frac, grid_dims, radius_per_atom, dtype): + """Anchor index, Cartesian sub-voxel residual ``w0``, and radius buckets.""" + device = xyz.device nx, ny, nz = grid_dims - ny_nz = ny * nz - strides = torch.tensor([ny_nz, nz, 1], device=device, dtype=torch.long) - for box_radius, b0, b1 in spans: - offsets = _get_radius_offsets(voxel_size, box_radius * min_voxel, device) # (R,3) - for s in range(b0, b1, chunk): - e = min(s + chunk, b1) - vi = (center_idx[s:e].unsqueeze(1) + offsets.unsqueeze(0)) % grid_shape - voxel_frac = vi.to(dtype) * inv_grid - diff = voxel_frac - xyz_frac[s:e].unsqueeze(1) - diff = diff - torch.round(diff) - r_sq = torch.einsum("avi,ij,avj->av", diff, G, diff) - expo = -_PI_SQ * r_sq.unsqueeze(2) / B_total[s:e].unsqueeze(1) - dens = torch.einsum("ag,avg->av", A_norm[s:e], torch.exp(expo)) - idx_flat = (vi.to(torch.long) * strides).sum(-1).view(-1) - density_flat = density_flat.scatter_add(0, idx_flat, dens.reshape(-1)) - return density_flat + grid_f = torch.tensor(grid_dims, device=device, dtype=dtype) + xyz_frac = (xyz @ inv_frac.T) % 1.0 + center_idx = torch.round(xyz_frac * grid_f).to(torch.long) + # w0: atom position relative to its anchor node, in Cartesian. This is what + # centres the sphere on the atom rather than on the node. + w0 = (xyz_frac - center_idx.to(dtype) / grid_f) @ frac.T + center_1d = ((center_idx[:, 0] % nx) * (ny * nz) + + (center_idx[:, 1] % ny) * nz + (center_idx[:, 2] % nz)) + order, spans = _bucket_by_radius(radius_per_atom, center_1d) + return order, spans, center_idx[order], w0[order] def add_isotropic_plain_var(density_map, xyz, adp, occ, A, B, - inv_frac_matrix, frac_matrix, grid_shape_tuple, - voxel_size, radius_per_atom): - """Variable-radius grouped fused-sphere isotropic splat; adds into ``density_map``. - - Uses only plain ``scatter_add`` (no custom C++ scatter), so it runs on every - device (CPU/CUDA/MPS), supports float64, and is double-differentiable -- the - portable / eager per-atom path.""" - nx, ny, nz = grid_shape_tuple - st = _iso_setup(xyz, adp, occ, A, B, inv_frac_matrix, frac_matrix, - grid_shape_tuple, voxel_size, radius_per_atom) + inv_frac_matrix, frac_matrix, radius_per_atom): + """Portable canonical-sphere isotropic splat; adds into ``density_map``. + + Signature mirrors + :func:`~torchref.base.electron_density.kernels.cpu.sphere_splat.add_isotropic_cpu_sphere_var` + exactly, so the two are interchangeable at the dispatch site. There is no + ``voxel_size`` or ``grid_shape`` argument: the grid shape comes from + ``density_map`` and the truncation box from ``inv_frac_matrix``. + """ + device, dtype = xyz.device, density_map.dtype + nx, ny, nz = (int(s) for s in density_map.shape) + grid_dims = (nx, ny, nz) + strides = torch.tensor([ny * nz, nz, 1], device=device, dtype=torch.long) + grid_shape = torch.tensor(grid_dims, device=device, dtype=torch.long) + + order, spans, center_idx, w0 = _canonical_setup( + xyz, inv_frac_matrix, frac_matrix, grid_dims, radius_per_atom, dtype) + B_total = ((B + adp[:, None]) * 0.25).clamp(min=0.1)[order] + A_norm = (A * occ[:, None])[order] * _PI_1P5 / (B_total * torch.sqrt(B_total)) + r2cut = (radius_per_atom * radius_per_atom)[order] + density_flat = density_map.reshape(-1) - density_flat = _splat_chunked_sphere( - density_flat, st["G"], st["inv_grid"], st["grid_dims"], voxel_size, - xyz.device, xyz.dtype, st["xyz_frac"], st["center_idx"], - st["B_total"], st["A_norm"], st["spans"], - ) + for r, b0, b1 in spans: + bh = _axis_half_widths(r, inv_frac_matrix, grid_dims) + offsets, off_cart = _box_offsets(bh, frac_matrix, grid_dims, device, dtype) + for s in range(b0, b1, _CHUNK): + e = min(s + _CHUNK, b1) + w = off_cart.unsqueeze(0) - w0[s:e].unsqueeze(1) # (C,R,3) + r_sq = (w * w).sum(-1) # (C,R) + expo = -_PI_SQ * r_sq.unsqueeze(2) / B_total[s:e].unsqueeze(1) + dens = torch.einsum("ag,avg->av", A_norm[s:e], torch.exp(expo)) + dens = torch.where(r_sq <= r2cut[s:e, None], dens, dens.new_zeros(())) + vi = (center_idx[s:e].unsqueeze(1) + offsets.unsqueeze(0)) % grid_shape + idx_flat = (vi * strides).sum(-1).reshape(-1) + density_flat = density_flat.scatter_add(0, idx_flat, dens.reshape(-1)) return density_flat.view(nx, ny, nz) -# ========================================================================= -# Anisotropic box-splat (full 3D Gaussian, no factorization) -# ========================================================================= -def _splat_chunked_aniso(density_flat, map_size, frac, inv_grid, grid_dims, device, dtype, - xyz_frac, center_idx, Minv, A_norm, spans, chunk=_CHUNK): - """Per box-bucket x chunk: full 3D aniso cube (`_aniso_density_cube`) + scatter.""" - nx, ny, nz = grid_dims - ny_nz = ny * nz - for box_radius, b0, b1 in spans: - axis = _axis_offsets(box_radius, device, dtypes.int) - axis_frac = axis.to(dtype).unsqueeze(0) * inv_grid.unsqueeze(1) - for s in range(b0, b1, chunk): - e = min(s + chunk, b1) - ci = center_idx[s:e] - sub = xyz_frac[s:e] - ci.to(dtype) * inv_grid - d_frac = axis_frac.unsqueeze(0) - sub.unsqueeze(2) - d_frac = d_frac - torch.round(d_frac) - cube = _aniso_density_cube(d_frac, frac, Minv[s:e], A_norm[s:e]) - wa = (ci[:, 0:1] + axis.unsqueeze(0)) % nx * ny_nz - wb = (ci[:, 1:2] + axis.unsqueeze(0)) % ny * nz - wc = (ci[:, 2:3] + axis.unsqueeze(0)) % nz - wbwc = wb.unsqueeze(2) + wc.unsqueeze(1) - density_flat = _do_structured_scatter(cube, wa, wbwc, density_flat, map_size) - return density_flat - - -def _aniso_setup(xyz, u, occ, A, B, inv_frac, frac, grid_shape, voxel_size, - radius_per_atom): - """Per-atom aniso M/Minv/A_norm + box-size buckets.""" - device, dtype = xyz.device, xyz.dtype - nx, ny, nz = grid_shape - ny_nz = ny * nz - grid_f = torch.tensor(grid_shape, device=device, dtype=dtype) - inv_grid = 1.0 / grid_f - N = xyz.shape[0] - - xyz_frac = xyz @ inv_frac.T - center_idx = torch.round((xyz_frac % 1.0) * grid_f).to(dtypes.int) - eye = torch.eye(3, dtype=dtype, device=device) - U3 = _u6_to_u3(u) - M = (B[:, :, None, None] * eye + _EIGHT_PI_SQ * U3[:, None, :, :]) / 4.0 # (N,5,3,3) - Minv = torch.linalg.inv(M) - det = torch.linalg.det(M).clamp(min=1e-10) - A_norm = A * occ[:, None] * _PI_1P5 / torch.sqrt(det) # (N,5) - - box_radius = _box_radius_per_atom(radius_per_atom, voxel_size) - center_1d = (center_idx[:, 0].long() * ny_nz - + center_idx[:, 1].long() * nz + center_idx[:, 2].long()) - order, spans = _bucket_by_box(box_radius, center_1d) - return dict( - frac=frac, inv_grid=inv_grid, grid_dims=(nx, ny, nz), - xyz_frac=xyz_frac[order], center_idx=center_idx[order], - Minv=Minv[order], A_norm=A_norm[order], spans=spans, - ) - +def add_anisotropic_plain_var(density_map, xyz, u, occ, A, B, + inv_frac_matrix, frac_matrix, radius_per_atom): + """Portable canonical-sphere anisotropic splat; adds into ``density_map``. -def add_anisotropic_cpu_var(real_space_grid, density_map, xyz, u, occ, A, B, - inv_frac_matrix, frac_matrix, radius_per_atom, voxel_size): - """Variable-radius grouped anisotropic box-splat; adds into ``density_map``.""" - nx, ny, nz = real_space_grid.shape[:3] - map_size = nx * ny * nz - st = _aniso_setup(xyz, u, occ, A, B, inv_frac_matrix, frac_matrix, - (nx, ny, nz), voxel_size, radius_per_atom) - density_flat = density_map.reshape(-1) - density_flat = _splat_chunked_aniso( - density_flat, map_size, st["frac"], st["inv_grid"], st["grid_dims"], - xyz.device, xyz.dtype, st["xyz_frac"], st["center_idx"], - st["Minv"], st["A_norm"], st["spans"], - ) - return density_flat.view(nx, ny, nz) + Signature mirrors + :func:`~torchref.base.electron_density.kernels.cpu.sphere_splat.add_anisotropic_cpu_sphere_var`. + Density is the full 3D Gaussian ``exp(-pi^2 w^T Minv w)`` with + ``M_g = (B_g*I + 8*pi^2*U)/4``; the *cutoff* is the Euclidean sphere at + ``radius_per_atom`` (the ellipsoid's isotropic bounding radius), matching the + CUDA and Metal kernels, which likewise cull on Euclidean distance and evaluate + the Mahalanobis form. + This replaces a full-cube splat: the cube was ~2.3x the sphere's voxel count + and disagreed with every accelerator path. + """ + device, dtype = xyz.device, density_map.dtype + nx, ny, nz = (int(s) for s in density_map.shape) + grid_dims = (nx, ny, nz) + strides = torch.tensor([ny * nz, nz, 1], device=device, dtype=torch.long) + grid_shape = torch.tensor(grid_dims, device=device, dtype=torch.long) + + order, spans, center_idx, w0 = _canonical_setup( + xyz, inv_frac_matrix, frac_matrix, grid_dims, radius_per_atom, dtype) + eye = torch.eye(3, dtype=dtype, device=device) + M = (B[:, :, None, None] * eye + _EIGHT_PI_SQ * _u6_to_u3(u)[:, None, :, :]) / 4.0 + Minv = torch.linalg.inv(M)[order] # (N,5,3,3) + det = torch.linalg.det(M).clamp(min=1e-10)[order] + A_norm = (A * occ[:, None])[order] * _PI_1P5 / torch.sqrt(det) # (N,5) + r2cut = (radius_per_atom * radius_per_atom)[order] -def _splat_chunked_aniso_plain(density_flat, frac, inv_grid, grid_dims, device, dtype, - xyz_frac, center_idx, Minv, A_norm, spans, chunk=_CHUNK): - """Per box-bucket x chunk: full 3D aniso cube + plain ``scatter_add`` (portable).""" - nx, ny, nz = grid_dims - ny_nz = ny * nz - for box_radius, b0, b1 in spans: - axis = _axis_offsets(box_radius, device, dtypes.int) - axis_frac = axis.to(dtype).unsqueeze(0) * inv_grid.unsqueeze(1) - for s in range(b0, b1, chunk): - e = min(s + chunk, b1) - ci = center_idx[s:e] - sub = xyz_frac[s:e] - ci.to(dtype) * inv_grid - d_frac = axis_frac.unsqueeze(0) - sub.unsqueeze(2) - d_frac = d_frac - torch.round(d_frac) - cube = _aniso_density_cube(d_frac, frac, Minv[s:e], A_norm[s:e]) # (C,L,L,L) - ix = (ci[:, 0:1] + axis.unsqueeze(0)) % nx # (C,L) - iy = (ci[:, 1:2] + axis.unsqueeze(0)) % ny - iz = (ci[:, 2:3] + axis.unsqueeze(0)) % nz - flat = (ix[:, :, None, None].to(torch.long) * ny_nz - + iy[:, None, :, None].to(torch.long) * nz - + iz[:, None, None, :].to(torch.long)) # (C,L,L,L) - density_flat = density_flat.scatter_add(0, flat.reshape(-1), cube.reshape(-1)) - return density_flat - - -def add_anisotropic_plain_var(real_space_grid, density_map, xyz, u, occ, A, B, - inv_frac_matrix, frac_matrix, radius_per_atom, voxel_size): - """Variable-radius anisotropic box-splat using plain ``scatter_add``. - - The portable / eager anisotropic per-atom path: no custom C++ scatter, so it - runs on every device, supports float64, and is double-differentiable.""" - nx, ny, nz = real_space_grid.shape[:3] - st = _aniso_setup(xyz, u, occ, A, B, inv_frac_matrix, frac_matrix, - (nx, ny, nz), voxel_size, radius_per_atom) density_flat = density_map.reshape(-1) - density_flat = _splat_chunked_aniso_plain( - density_flat, st["frac"], st["inv_grid"], st["grid_dims"], - xyz.device, xyz.dtype, st["xyz_frac"], st["center_idx"], - st["Minv"], st["A_norm"], st["spans"], - ) + for r, b0, b1 in spans: + bh = _axis_half_widths(r, inv_frac_matrix, grid_dims) + offsets, off_cart = _box_offsets(bh, frac_matrix, grid_dims, device, dtype) + for s in range(b0, b1, _CHUNK): + e = min(s + _CHUNK, b1) + w = off_cart.unsqueeze(0) - w0[s:e].unsqueeze(1) # (C,R,3) + r_sq = (w * w).sum(-1) # (C,R) + # q[a,v,g] = w^T Minv_g w + q = torch.einsum("avi,agij,avj->avg", w, Minv[s:e], w) + dens = torch.einsum("ag,avg->av", A_norm[s:e], torch.exp(-_PI_SQ * q)) + dens = torch.where(r_sq <= r2cut[s:e, None], dens, dens.new_zeros(())) + vi = (center_idx[s:e].unsqueeze(1) + offsets.unsqueeze(0)) % grid_shape + idx_flat = (vi * strides).sum(-1).reshape(-1) + density_flat = density_flat.scatter_add(0, idx_flat, dens.reshape(-1)) return density_flat.view(nx, ny, nz) diff --git a/torchref/base/electron_density/kernels/cuda/variable_radius.py b/torchref/base/electron_density/kernels/cuda/variable_radius.py index 8c9ba5ab..755757c6 100644 --- a/torchref/base/electron_density/kernels/cuda/variable_radius.py +++ b/torchref/base/electron_density/kernels/cuda/variable_radius.py @@ -811,3 +811,106 @@ def backward(ctx, grad_density_map): # out = density_map + splat -> grad wrt density_map is identity. return (grad_density_map, None, grad_xyz, grad_u, grad_occ, None, None, None, None, None, None) + + +# --------------------------------------------------------------------------- +# Canonical-signature wrappers +# --------------------------------------------------------------------------- +# Every production splat exposes the same signature: +# +# (density_map, xyz, adp_or_u, occ, A, B, inv_frac_matrix, frac_matrix, +# radius_per_atom) +# +# so the dispatch in ``electron_density/main.py`` is four structurally identical +# calls per ladder and a test can drive any backend through one code path. These +# wrappers do for CUDA what ``add_*_mps_var`` already did for Metal: square the +# radius and build the coefficient mask, rather than leaving that to the caller. +# +# ``density_map`` is passed where the ``autograd.Function`` wants +# ``real_space_grid``. That is exact, not a convenience: ``forward``/``backward`` +# use that argument only for ``.shape[:3]`` and for a ``grid_flat`` pointer handed +# to the kernel as ``grid_ptr`` -- which none of the four Triton kernels ever +# loads, because voxel coordinates are derived arithmetically from ``frac`` and the +# grid dims. ``density_map`` has the same ``(nx, ny, nz)`` shape, so both uses are +# satisfied. If a kernel is ever changed to actually read ``grid_ptr``, this breaks +# and the right fix is to delete that dead parameter, not to thread a real grid +# back through here. + + +def why_unavailable(): + """``None`` if these kernels can run, else why they cannot. + + The reason-returning half of the availability protocol shared by every backend (see + :mod:`torchref.utils.backends`). + + This probe closes a real gap rather than restating ``triton_available()``. The + ``@triton.jit`` kernel bodies live inside ``if _HAVE_TRITON:`` above, but + ``WorkQueueGridDensity`` and these wrappers are defined *unconditionally*, so on a host + without Triton the module imports cleanly and the failure surfaces as a bare + ``NameError`` from deep inside ``_launch_grid_fwd``. Every other backend answers + "can I run" before being called; this one had nothing to ask. + """ + if not _HAVE_TRITON: + return ( + "triton is not importable, so the work-queue kernel bodies were never " + "compiled into this module" + ) + return None + + +def _coeff_mask(xyz): + """All-ones per-atom coefficient mask, shape ``(n, 5)``. + + Every call site in the codebase passes all ones -- the mask exists so a caller + *could* disable individual ITC92 Gaussians, and nothing does. Kept because it is + part of the kernel's argument list. + """ + return torch.ones(xyz.shape[0], 5, dtype=xyz.dtype, device=xyz.device) + + +def add_isotropic_cuda_var( + density_map, xyz, adp, occ, A, B, inv_frac_matrix, frac_matrix, radius_per_atom +): + """Isotropic variable-radius Triton splat; returns ``density_map + splat``. + + Canonical splat signature, identical to ``add_isotropic_plain_var``, + ``add_isotropic_cpu_sphere_var`` and ``add_isotropic_mps_var``. CUDA float32 + only -- the gate is ``should_use_triton``; this wrapper does not re-check. + """ + return WorkQueueGridDensity.apply( + density_map, + density_map, # stands in for real_space_grid; see the note above + xyz, + adp, + occ, + A, + B, + radius_per_atom * radius_per_atom, + _coeff_mask(xyz), + inv_frac_matrix, + frac_matrix, + ) + + +def add_anisotropic_cuda_var( + density_map, xyz, u, occ, A, B, inv_frac_matrix, frac_matrix, radius_per_atom +): + """Anisotropic variable-radius Triton splat; returns ``density_map + splat``. + + Canonical splat signature. ``u`` carries ``[U11, U22, U33, U12, U13, U23]``; the + cutoff stays the Euclidean sphere at ``radius_per_atom`` while the density is the + full Mahalanobis form, matching the Metal and fused-CPU kernels. + """ + return WorkQueueGridDensityAniso.apply( + density_map, + density_map, # stands in for real_space_grid; see the note above + xyz, + u, + occ, + A, + B, + radius_per_atom * radius_per_atom, + _coeff_mask(xyz), + inv_frac_matrix, + frac_matrix, + ) diff --git a/torchref/base/electron_density/kernels/mps/__init__.py b/torchref/base/electron_density/kernels/mps/__init__.py new file mode 100644 index 00000000..0ffdc5de --- /dev/null +++ b/torchref/base/electron_density/kernels/mps/__init__.py @@ -0,0 +1,29 @@ +"""Metal (MPS) variable-radius electron-density splat kernels. + +Native Metal kernels (compiled at runtime via ``torch.mps.compile_shader``) for +the density splat on Apple-silicon GPUs, replacing the portable eager splat that +dominates fcalc time on MPS. Selection goes through +``torchref.utils.should_use_metal``, which gates on MPS + float32 + a compiled +shader; a runtime failure degrades to the portable splat and warns. +Every other platform is unaffected. +""" + +from torchref.base.electron_density.kernels.mps.compile import ( + clear_cache, + last_error, + mps_kernels_available, + warmup, +) +from torchref.base.electron_density.kernels.mps.variable_radius import ( + add_anisotropic_mps_var, + add_isotropic_mps_var, +) + +__all__ = [ + "add_isotropic_mps_var", + "add_anisotropic_mps_var", + "mps_kernels_available", + "warmup", + "clear_cache", + "last_error", +] diff --git a/torchref/base/electron_density/kernels/mps/_shaders.py b/torchref/base/electron_density/kernels/mps/_shaders.py new file mode 100644 index 00000000..facd694c --- /dev/null +++ b/torchref/base/electron_density/kernels/mps/_shaders.py @@ -0,0 +1,434 @@ +"""Metal Shading Language (MSL) source for the variable-radius density splat. + +Compiled at runtime via ``torch.mps.compile_shader`` (see ``compile.py``). One +GPU thread per atom (dispatched ``threads=[n_atoms]``): each thread sizes its +own cubic bounding box from ``r2cut`` and the inverse-cell row norms, iterates +the box, truncates to the per-atom sphere (``r2 <= r2cut``), evaluates the +5-term ITC92 Gaussian, and accumulates into the grid via a portable +compare-exchange float atomic-add (``atomic_add_f``) that works on every Metal +GPU family (Apple7 / M1 onward), not just Metal-3 native ``atomic_float``. + +The math mirrors the CUDA variable-radius kernels +(``kernels/cuda/variable_radius.py``) and is validated against the portable CPU +reference (``add_isotropic_plain_var`` / ``add_anisotropic_plain_var``). Because +each thread owns one atom, the backward kernels accumulate per-atom gradients +with no atomics. + +Coordinates: work in Cartesian offsets. ``frac`` is the fractional->Cartesian +matrix (its columns are the cell vectors a,b,c); ``inv_frac`` is its inverse +(Cartesian->fractional). For voxel offset ``o`` from the atom's grid anchor, +``w = frac @ (o/n - residual)`` is the Cartesian atom->voxel vector and +``r2 = w.w``. + +Constants match the reference bit-for-bit: PI_1P5 = pi^1.5, PI_SQ = pi^2. +""" + +# NOTE: kept as a single translation unit so one compile_shader call yields all +# kernels. Backward (iso/aniso) and anisotropic kernels are appended below as +# each is validated. +MSL_SOURCE = r""" +#include +#include +using namespace metal; + +constant float PI_1P5 = 5.568327996831708f; // pi^1.5 +constant float PI_SQ = 9.869604401089358f; // pi^2 + +// Portable float atomic-add via compare-exchange on uint. Works on EVERY Metal +// GPU family (incl. Apple7 / M1), not just Metal-3 native atomic_float. Measured +// identical to native on Apple8/M2 -- the splat's per-atom locality keeps +// contention low, so the CAS loop essentially never retries. +inline void atomic_add_f(device atomic_uint* addr, float val) { + uint old = atomic_load_explicit(addr, memory_order_relaxed), nxt; + do { nxt = as_type(as_type(old) + val); } + while (!atomic_compare_exchange_weak_explicit( + addr, &old, nxt, memory_order_relaxed, memory_order_relaxed)); +} + +// --------------------------------------------------------------------------- +// Isotropic forward: accumulate 5-Gaussian density into grid (atomic add). +// --------------------------------------------------------------------------- +kernel void iso_splat_fwd( + device atomic_uint* grid [[buffer(0)]], // (nx*ny*nz,) accumulator + device const float* xyz [[buffer(1)]], // (n,3) Cartesian + device const float* adp [[buffer(2)]], // (n,) isotropic B + device const float* occ [[buffer(3)]], // (n,) + device const float* A [[buffer(4)]], // (n,5) ITC92 amplitudes + device const float* B [[buffer(5)]], // (n,5) ITC92 widths + device const float* r2cut [[buffer(6)]], // (n,) squared radius + device const float* mask [[buffer(7)]], // (n,5) per-Gaussian mask + device const float* inv_frac [[buffer(8)]], // 9 (row-major 3x3) + device const float* frac [[buffer(9)]], // 9 (row-major 3x3) + constant int& n_atoms [[buffer(10)]], + constant int& nx [[buffer(11)]], + constant int& ny [[buffer(12)]], + constant int& nz [[buffer(13)]], + uint a [[thread_position_in_grid]]) +{ + if (a >= (uint)n_atoms) return; + + // Cell geometry (frac columns are cell vectors a,b,c). + float f0=frac[0],f1=frac[1],f2=frac[2], + f3=frac[3],f4=frac[4],f5=frac[5], + f6=frac[6],f7=frac[7],f8=frac[8]; + float i0=inv_frac[0],i1=inv_frac[1],i2=inv_frac[2], + i3=inv_frac[3],i4=inv_frac[4],i5=inv_frac[5], + i6=inv_frac[6],i7=inv_frac[7],i8=inv_frac[8]; + float fnx=(float)nx, fny=(float)ny, fnz=(float)nz; + + // Per-axis Cartesian voxel step vectors (cell vector / n). + float uax=f0/fnx, uay=f3/fnx, uaz=f6/fnx; + float ubx=f1/fny, uby=f4/fny, ubz=f7/fny; + float ucx=f2/fnz, ucy=f5/fnz, ucz=f8/fnz; + + float rc2 = r2cut[a]; + float r = sqrt(rc2); + // Triclinic-correct half-widths: max|off_axis| = n_axis * r * ||inv_frac row||. + int bhx = (int)ceil(r*fnx*sqrt(i0*i0+i1*i1+i2*i2)); + int bhy = (int)ceil(r*fny*sqrt(i3*i3+i4*i4+i5*i5)); + int bhz = (int)ceil(r*fnz*sqrt(i6*i6+i7*i7+i8*i8)); + + // Atom -> fractional -> wrap into [0,1) -> nearest grid node + residual. + float ax=xyz[3*a+0], ay=xyz[3*a+1], az=xyz[3*a+2]; + float fx=ax*i0+ay*i1+az*i2; + float fy=ax*i3+ay*i4+az*i5; + float fz=ax*i6+ay*i7+az*i8; + fx-=floor(fx); fy-=floor(fy); fz-=floor(fz); + int cix=(int)rint(fx*fnx), ciy=(int)rint(fy*fny), ciz=(int)rint(fz*fnz); + float sx=fx-(float)cix/fnx, sy=fy-(float)ciy/fny, sz=fz-(float)ciz/fnz; + float w0x=f0*sx+f1*sy+f2*sz; // residual -> Cartesian + float w0y=f3*sx+f4*sy+f5*sz; + float w0z=f6*sx+f7*sy+f8*sz; + + // Per-Gaussian amplitude / width. + float occa=occ[a], b_iso=adp[a]; + float Bt[5], An[5]; + for (int g=0; g<5; ++g) { + float Bt_g = max((B[5*a+g]+b_iso)*0.25f, 0.1f); + Bt[g] = Bt_g; + An[g] = mask[5*a+g]*A[5*a+g]*occa*PI_1P5/(Bt_g*sqrt(Bt_g)); + } + + for (int ox=-bhx; ox<=bhx; ++ox) { + float fox=(float)ox; + for (int oy=-bhy; oy<=bhy; ++oy) { + float foy=(float)oy; + for (int oz=-bhz; oz<=bhz; ++oz) { + float foz=(float)oz; + float wx=fox*uax+foy*ubx+foz*ucx - w0x; + float wy=fox*uay+foy*uby+foz*ucy - w0y; + float wz=fox*uaz+foy*ubz+foz*ucz - w0z; + float r2=wx*wx+wy*wy+wz*wz; + if (r2 > rc2) continue; + float dens=0.0f; + for (int g=0; g<5; ++g) dens += An[g]*fast::exp(-PI_SQ*r2/Bt[g]); + int vix=cix+ox; vix=((vix%nx)+nx)%nx; + int viy=ciy+oy; viy=((viy%ny)+ny)%ny; + int viz=ciz+oz; viz=((viz%nz)+nz)%nz; + int idx=(vix*ny+viy)*nz+viz; + atomic_add_f(&grid[idx], dens); + } + } + } +} + +// --------------------------------------------------------------------------- +// Isotropic backward: analytic grads to xyz, adp, occ. One thread per atom, so +// each thread owns its atom's gradient slots -> no atomics. Recomputes the +// forward geometry and gathers the upstream grad from the grid. +// --------------------------------------------------------------------------- +kernel void iso_splat_bwd( + device float* grad_xyz [[buffer(0)]], // (n,3) out + device float* grad_adp [[buffer(1)]], // (n,) out + device float* grad_occ [[buffer(2)]], // (n,) out + device const float* grad_out [[buffer(3)]], // (nx*ny*nz,) upstream grad + device const float* xyz [[buffer(4)]], + device const float* adp [[buffer(5)]], + device const float* occ [[buffer(6)]], + device const float* A [[buffer(7)]], + device const float* B [[buffer(8)]], + device const float* r2cut [[buffer(9)]], + device const float* mask [[buffer(10)]], + device const float* inv_frac [[buffer(11)]], + device const float* frac [[buffer(12)]], + constant int& n_atoms [[buffer(13)]], + constant int& nx [[buffer(14)]], + constant int& ny [[buffer(15)]], + constant int& nz [[buffer(16)]], + uint a [[thread_position_in_grid]]) +{ + if (a >= (uint)n_atoms) return; + + float f0=frac[0],f1=frac[1],f2=frac[2], + f3=frac[3],f4=frac[4],f5=frac[5], + f6=frac[6],f7=frac[7],f8=frac[8]; + float i0=inv_frac[0],i1=inv_frac[1],i2=inv_frac[2], + i3=inv_frac[3],i4=inv_frac[4],i5=inv_frac[5], + i6=inv_frac[6],i7=inv_frac[7],i8=inv_frac[8]; + float fnx=(float)nx, fny=(float)ny, fnz=(float)nz; + float uax=f0/fnx, uay=f3/fnx, uaz=f6/fnx; + float ubx=f1/fny, uby=f4/fny, ubz=f7/fny; + float ucx=f2/fnz, ucy=f5/fnz, ucz=f8/fnz; + + float rc2 = r2cut[a]; + float r = sqrt(rc2); + int bhx = (int)ceil(r*fnx*sqrt(i0*i0+i1*i1+i2*i2)); + int bhy = (int)ceil(r*fny*sqrt(i3*i3+i4*i4+i5*i5)); + int bhz = (int)ceil(r*fnz*sqrt(i6*i6+i7*i7+i8*i8)); + + float ax=xyz[3*a+0], ay=xyz[3*a+1], az=xyz[3*a+2]; + float fx=ax*i0+ay*i1+az*i2; + float fy=ax*i3+ay*i4+az*i5; + float fz=ax*i6+ay*i7+az*i8; + fx-=floor(fx); fy-=floor(fy); fz-=floor(fz); + int cix=(int)rint(fx*fnx), ciy=(int)rint(fy*fny), ciz=(int)rint(fz*fnz); + float sx=fx-(float)cix/fnx, sy=fy-(float)ciy/fny, sz=fz-(float)ciz/fnz; + float w0x=f0*sx+f1*sy+f2*sz; + float w0y=f3*sx+f4*sy+f5*sz; + float w0z=f6*sx+f7*sy+f8*sz; + + float occa=occ[a], b_iso=adp[a]; + float Bt[5], An[5], clampf[5]; + for (int g=0; g<5; ++g) { + float raw = (B[5*a+g]+b_iso)*0.25f; + float Bt_g = max(raw, 0.1f); + Bt[g] = Bt_g; + An[g] = mask[5*a+g]*A[5*a+g]*occa*PI_1P5/(Bt_g*sqrt(Bt_g)); + clampf[g] = (raw > 0.1f) ? 1.0f : 0.0f; // d(clamped Bt)/d(adp)=0 in clamp region + } + + float gx=0.0f, gy=0.0f, gz=0.0f, gb=0.0f, go=0.0f; + for (int ox=-bhx; ox<=bhx; ++ox) { + float fox=(float)ox; + for (int oy=-bhy; oy<=bhy; ++oy) { + float foy=(float)oy; + for (int oz=-bhz; oz<=bhz; ++oz) { + float foz=(float)oz; + float wx=fox*uax+foy*ubx+foz*ucx - w0x; + float wy=fox*uay+foy*uby+foz*ucy - w0y; + float wz=fox*uaz+foy*ubz+foz*ucz - w0z; + float r2=wx*wx+wy*wy+wz*wz; + if (r2 > rc2) continue; + int vix=cix+ox; vix=((vix%nx)+nx)%nx; + int viy=ciy+oy; viy=((viy%ny)+ny)%ny; + int viz=ciz+oz; viz=((viz%nz)+nz)%nz; + float g_out = grad_out[(vix*ny+viy)*nz+viz]; + float dens=0.0f, coeff_xyz=0.0f, db_sum=0.0f; + for (int g=0; g<5; ++g) { + float e = fast::exp(-PI_SQ*r2/Bt[g]); + float Ae = An[g]*e; + dens += Ae; + coeff_xyz += Ae/Bt[g]; + db_sum += Ae*(-1.5f/Bt[g] + PI_SQ*r2/(Bt[g]*Bt[g]))*clampf[g]; + } + float sxyz = g_out*2.0f*PI_SQ*coeff_xyz; + gx += sxyz*wx; gy += sxyz*wy; gz += sxyz*wz; + gb += g_out*0.25f*db_sum; + go += g_out*dens; + } + } + } + grad_xyz[3*a+0]=gx; grad_xyz[3*a+1]=gy; grad_xyz[3*a+2]=gz; + grad_adp[a]=gb; + grad_occ[a]=(occa!=0.0f) ? go/occa : 0.0f; +} + +constant float TWO_PI_SQ = 19.739208802178716f; // 2*pi^2 (= 8*pi^2 / 4) + +// --------------------------------------------------------------------------- +// Anisotropic forward. Per Gaussian g: M_g = (B_g*I + 8*pi^2*U)/4 (symmetric +// 3x3), inverted analytically; A_norm_g = mask*A*occ*pi^1.5/sqrt(det M_g); +// q_g = w^T Minv_g w; density = sum_g A_norm_g fast::exp(-pi^2 q_g). Truncation +// is a per-axis bounding box (from r2cut + inv-cell row norms) with a sphere +// cull (r2 <= r2cut) -- far tighter than a full cube on anisotropic high-res +// cells. Uses metal::fast::exp (accurate to ~1e-4, well under the grid floor). +// --------------------------------------------------------------------------- +kernel void aniso_splat_fwd( + device atomic_uint* grid [[buffer(0)]], + device const float* xyz [[buffer(1)]], // (n,3) + device const float* u [[buffer(2)]], // (n,6) U11,U22,U33,U12,U13,U23 + device const float* occ [[buffer(3)]], + device const float* A [[buffer(4)]], // (n,5) + device const float* B [[buffer(5)]], // (n,5) + device const float* r2cut [[buffer(6)]], // (n,) squared truncation radius + device const float* mask [[buffer(7)]], // (n,5) + device const float* inv_frac [[buffer(8)]], + device const float* frac [[buffer(9)]], + constant int& n_atoms [[buffer(10)]], + constant int& nx [[buffer(11)]], + constant int& ny [[buffer(12)]], + constant int& nz [[buffer(13)]], + uint a [[thread_position_in_grid]]) +{ + if (a >= (uint)n_atoms) return; + + float f0=frac[0],f1=frac[1],f2=frac[2], + f3=frac[3],f4=frac[4],f5=frac[5], + f6=frac[6],f7=frac[7],f8=frac[8]; + float i0=inv_frac[0],i1=inv_frac[1],i2=inv_frac[2], + i3=inv_frac[3],i4=inv_frac[4],i5=inv_frac[5], + i6=inv_frac[6],i7=inv_frac[7],i8=inv_frac[8]; + float fnx=(float)nx, fny=(float)ny, fnz=(float)nz; + float uax=f0/fnx, uay=f3/fnx, uaz=f6/fnx; + float ubx=f1/fny, uby=f4/fny, ubz=f7/fny; + float ucx=f2/fnz, ucy=f5/fnz, ucz=f8/fnz; + + float ax=xyz[3*a+0], ay=xyz[3*a+1], az=xyz[3*a+2]; + float fx=ax*i0+ay*i1+az*i2, fy=ax*i3+ay*i4+az*i5, fz=ax*i6+ay*i7+az*i8; + fx-=floor(fx); fy-=floor(fy); fz-=floor(fz); + int cix=(int)rint(fx*fnx), ciy=(int)rint(fy*fny), ciz=(int)rint(fz*fnz); + float sx=fx-(float)cix/fnx, sy=fy-(float)ciy/fny, sz=fz-(float)ciz/fnz; + float w0x=f0*sx+f1*sy+f2*sz, w0y=f3*sx+f4*sy+f5*sz, w0z=f6*sx+f7*sy+f8*sz; + + float occa=occ[a]; + float u11=u[6*a+0],u22=u[6*a+1],u33=u[6*a+2],u12=u[6*a+3],u13=u[6*a+4],u23=u[6*a+5]; + float p00[5],p11[5],p22[5],p01[5],p02[5],p12[5],An[5]; + for (int g=0; g<5; ++g) { + float Bg=B[5*a+g]; + float ma=0.25f*Bg+TWO_PI_SQ*u11, mb=0.25f*Bg+TWO_PI_SQ*u22, mc=0.25f*Bg+TWO_PI_SQ*u33; + float md=TWO_PI_SQ*u12, me=TWO_PI_SQ*u13, mf=TWO_PI_SQ*u23; + float det=ma*(mb*mc-mf*mf)-md*(md*mc-me*mf)+me*(md*mf-me*mb); + float inv=1.0f/det; + p00[g]=(mb*mc-mf*mf)*inv; p11[g]=(ma*mc-me*me)*inv; p22[g]=(ma*mb-md*md)*inv; + p01[g]=(me*mf-md*mc)*inv; p02[g]=(md*mf-me*mb)*inv; p12[g]=(md*me-ma*mf)*inv; + An[g]=mask[5*a+g]*A[5*a+g]*occa*PI_1P5/sqrt(max(det,1e-10f)); + } + + float rc2=r2cut[a], r=sqrt(rc2); + int bhx=(int)ceil(r*fnx*sqrt(i0*i0+i1*i1+i2*i2)); + int bhy=(int)ceil(r*fny*sqrt(i3*i3+i4*i4+i5*i5)); + int bhz=(int)ceil(r*fnz*sqrt(i6*i6+i7*i7+i8*i8)); + for (int ox=-bhx; ox<=bhx; ++ox) { + float fox=(float)ox; + for (int oy=-bhy; oy<=bhy; ++oy) { + float foy=(float)oy; + for (int oz=-bhz; oz<=bhz; ++oz) { + float foz=(float)oz; + float wx=fox*uax+foy*ubx+foz*ucx - w0x; + float wy=fox*uay+foy*uby+foz*ucy - w0y; + float wz=fox*uaz+foy*ubz+foz*ucz - w0z; + if (wx*wx+wy*wy+wz*wz > rc2) continue; + float dens=0.0f; + for (int g=0; g<5; ++g) { + float q=p00[g]*wx*wx+p11[g]*wy*wy+p22[g]*wz*wz + +2.0f*(p01[g]*wx*wy+p02[g]*wx*wz+p12[g]*wy*wz); + dens += An[g]*fast::exp(-PI_SQ*q); + } + int vix=cix+ox; vix=((vix%nx)+nx)%nx; + int viy=ciy+oy; viy=((viy%ny)+ny)%ny; + int viz=ciz+oz; viz=((viz%nz)+nz)%nz; + atomic_add_f(&grid[(vix*ny+viy)*nz+viz], dens); + } + } + } +} + +// --------------------------------------------------------------------------- +// Anisotropic backward: grads to xyz, U (6), occ. v_g = Minv_g w; dg_g = +// A_norm_g exp(-pi^2 w.v_g). grad_xyz = go*2pi^2*sum dg v; grad_U diag = +// go*2pi^2*sum dg(-0.5 p_ii + pi^2 v_i^2); grad_U offdiag = go*4pi^2*sum +// dg(-0.5 p_ij + pi^2 v_i v_j); grad_occ = go*sum dg / occ. One thread/atom. +// --------------------------------------------------------------------------- +kernel void aniso_splat_bwd( + device float* grad_xyz [[buffer(0)]], // (n,3) + device float* grad_u [[buffer(1)]], // (n,6) + device float* grad_occ [[buffer(2)]], // (n,) + device const float* grad_out [[buffer(3)]], + device const float* xyz [[buffer(4)]], + device const float* u [[buffer(5)]], + device const float* occ [[buffer(6)]], + device const float* A [[buffer(7)]], + device const float* B [[buffer(8)]], + device const float* r2cut [[buffer(9)]], + device const float* mask [[buffer(10)]], + device const float* inv_frac [[buffer(11)]], + device const float* frac [[buffer(12)]], + constant int& n_atoms [[buffer(13)]], + constant int& nx [[buffer(14)]], + constant int& ny [[buffer(15)]], + constant int& nz [[buffer(16)]], + uint a [[thread_position_in_grid]]) +{ + if (a >= (uint)n_atoms) return; + + float f0=frac[0],f1=frac[1],f2=frac[2], + f3=frac[3],f4=frac[4],f5=frac[5], + f6=frac[6],f7=frac[7],f8=frac[8]; + float i0=inv_frac[0],i1=inv_frac[1],i2=inv_frac[2], + i3=inv_frac[3],i4=inv_frac[4],i5=inv_frac[5], + i6=inv_frac[6],i7=inv_frac[7],i8=inv_frac[8]; + float fnx=(float)nx, fny=(float)ny, fnz=(float)nz; + float uax=f0/fnx, uay=f3/fnx, uaz=f6/fnx; + float ubx=f1/fny, uby=f4/fny, ubz=f7/fny; + float ucx=f2/fnz, ucy=f5/fnz, ucz=f8/fnz; + + float ax=xyz[3*a+0], ay=xyz[3*a+1], az=xyz[3*a+2]; + float fx=ax*i0+ay*i1+az*i2, fy=ax*i3+ay*i4+az*i5, fz=ax*i6+ay*i7+az*i8; + fx-=floor(fx); fy-=floor(fy); fz-=floor(fz); + int cix=(int)rint(fx*fnx), ciy=(int)rint(fy*fny), ciz=(int)rint(fz*fnz); + float sx=fx-(float)cix/fnx, sy=fy-(float)ciy/fny, sz=fz-(float)ciz/fnz; + float w0x=f0*sx+f1*sy+f2*sz, w0y=f3*sx+f4*sy+f5*sz, w0z=f6*sx+f7*sy+f8*sz; + + float occa=occ[a]; + float u11=u[6*a+0],u22=u[6*a+1],u33=u[6*a+2],u12=u[6*a+3],u13=u[6*a+4],u23=u[6*a+5]; + float p00[5],p11[5],p22[5],p01[5],p02[5],p12[5],An[5]; + for (int g=0; g<5; ++g) { + float Bg=B[5*a+g]; + float ma=0.25f*Bg+TWO_PI_SQ*u11, mb=0.25f*Bg+TWO_PI_SQ*u22, mc=0.25f*Bg+TWO_PI_SQ*u33; + float md=TWO_PI_SQ*u12, me=TWO_PI_SQ*u13, mf=TWO_PI_SQ*u23; + float det=ma*(mb*mc-mf*mf)-md*(md*mc-me*mf)+me*(md*mf-me*mb); + float inv=1.0f/det; + p00[g]=(mb*mc-mf*mf)*inv; p11[g]=(ma*mc-me*me)*inv; p22[g]=(ma*mb-md*md)*inv; + p01[g]=(me*mf-md*mc)*inv; p02[g]=(md*mf-me*mb)*inv; p12[g]=(md*me-ma*mf)*inv; + An[g]=mask[5*a+g]*A[5*a+g]*occa*PI_1P5/sqrt(max(det,1e-10f)); + } + + float gx=0,gy=0,gz=0,gu0=0,gu1=0,gu2=0,gu3=0,gu4=0,gu5=0,go=0; + float rc2=r2cut[a], r=sqrt(rc2); + int bhx=(int)ceil(r*fnx*sqrt(i0*i0+i1*i1+i2*i2)); + int bhy=(int)ceil(r*fny*sqrt(i3*i3+i4*i4+i5*i5)); + int bhz=(int)ceil(r*fnz*sqrt(i6*i6+i7*i7+i8*i8)); + for (int ox=-bhx; ox<=bhx; ++ox) { + float fox=(float)ox; + for (int oy=-bhy; oy<=bhy; ++oy) { + float foy=(float)oy; + for (int oz=-bhz; oz<=bhz; ++oz) { + float foz=(float)oz; + float wx=fox*uax+foy*ubx+foz*ucx - w0x; + float wy=fox*uay+foy*uby+foz*ucy - w0y; + float wz=fox*uaz+foy*ubz+foz*ucz - w0z; + if (wx*wx+wy*wy+wz*wz > rc2) continue; + int vix=cix+ox; vix=((vix%nx)+nx)%nx; + int viy=ciy+oy; viy=((viy%ny)+ny)%ny; + int viz=ciz+oz; viz=((viz%nz)+nz)%nz; + float g_out=grad_out[(vix*ny+viy)*nz+viz]; + float sxx=0,syy=0,szz=0,dens=0; + float su0=0,su1=0,su2=0,su3=0,su4=0,su5=0; + for (int g=0; g<5; ++g) { + float vx=p00[g]*wx+p01[g]*wy+p02[g]*wz; + float vy=p01[g]*wx+p11[g]*wy+p12[g]*wz; + float vz=p02[g]*wx+p12[g]*wy+p22[g]*wz; + float q=wx*vx+wy*vy+wz*vz; + float dg=An[g]*fast::exp(-PI_SQ*q); + dens+=dg; sxx+=dg*vx; syy+=dg*vy; szz+=dg*vz; + su0+=dg*(-0.5f*p00[g]+PI_SQ*vx*vx); + su1+=dg*(-0.5f*p11[g]+PI_SQ*vy*vy); + su2+=dg*(-0.5f*p22[g]+PI_SQ*vz*vz); + su3+=dg*(-0.5f*p01[g]+PI_SQ*vx*vy); + su4+=dg*(-0.5f*p02[g]+PI_SQ*vx*vz); + su5+=dg*(-0.5f*p12[g]+PI_SQ*vy*vz); + } + float s=g_out*2.0f*PI_SQ; + gx+=s*sxx; gy+=s*syy; gz+=s*szz; + gu0+=s*su0; gu1+=s*su1; gu2+=s*su2; + gu3+=g_out*4.0f*PI_SQ*su3; gu4+=g_out*4.0f*PI_SQ*su4; gu5+=g_out*4.0f*PI_SQ*su5; + go+=g_out*dens; + } + } + } + grad_xyz[3*a+0]=gx; grad_xyz[3*a+1]=gy; grad_xyz[3*a+2]=gz; + grad_u[6*a+0]=gu0; grad_u[6*a+1]=gu1; grad_u[6*a+2]=gu2; + grad_u[6*a+3]=gu3; grad_u[6*a+4]=gu4; grad_u[6*a+5]=gu5; + grad_occ[a]=(occa!=0.0f) ? go/occa : 0.0f; +} +""" diff --git a/torchref/base/electron_density/kernels/mps/compile.py b/torchref/base/electron_density/kernels/mps/compile.py new file mode 100644 index 00000000..3e13d8f0 --- /dev/null +++ b/torchref/base/electron_density/kernels/mps/compile.py @@ -0,0 +1,111 @@ +"""Lazy compilation + availability probe for the Metal (MPS) splat kernels. + +Mirrors the memoize / permanent-failure / graceful-fallback structure of +``kernels/cpu/scatter.py`` (``_get_module``), but the build is a single +in-process ``torch.mps.compile_shader`` call -- no ninja, build directory, or +file locking, since PyTorch caches the compiled pipeline-state objects itself. + +``mps_kernels_available()`` is what ``torchref.utils.should_use_metal`` consults; +it returns False whenever MPS is absent, ``compile_shader`` is missing +(torch < 2.9), or the shader fails to build. The caller then falls back to the +portable plain splat and warns; ``why_unavailable()`` reports +instead, quoting :func:`last_error`. +""" + +from __future__ import annotations + +import traceback +from typing import Optional, Tuple + +import torch + +from torchref.base.electron_density.kernels.mps._shaders import MSL_SOURCE + +# Memoized compile state (attempted at most once per process). +_lib = None +_lib_failed = False +_lib_error: Optional[Tuple[str, str]] = None # (message, traceback) + + +def _mps_shader_supported() -> bool: + """True iff this torch build can compile+run Metal shaders on this host.""" + return ( + hasattr(torch, "mps") + and hasattr(torch.mps, "compile_shader") + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ) + + +def _get_lib(): + """Return the compiled shader library, or None if unavailable. + + Short-circuits both prior outcomes so compilation is attempted exactly once. + Any failure (unsupported torch, MPS off, MSL compile error) is recorded and + turned into a None return so callers fall back to the plain splat. + """ + global _lib, _lib_failed, _lib_error + if _lib is not None: + return _lib + if _lib_failed: + return None + if not _mps_shader_supported(): + _lib_failed = True + _lib_error = ( + "torch.mps.compile_shader unavailable or MPS not available", + "", + ) + return None + try: + _lib = torch.mps.compile_shader(MSL_SOURCE) + except Exception as e: # noqa: BLE001 - any build failure -> fall back + _lib_failed = True + _lib_error = (f"{type(e).__name__}: {e}", traceback.format_exc()) + return None + return _lib + + +def why_unavailable() -> Optional[str]: + """``None`` if the Metal kernels are usable, else why they are not. + + The single availability probe for this backend -- the shape every backend implements, + consumed by :mod:`torchref.utils.backends`. It returns the *reason* rather than a bool + because the availability test has to explain itself, and "torch has no + ``compile_shader``" and "the MSL failed to compile" are different problems with + different fixes. + """ + if _get_lib() is not None: + return None + reason = _lib_error[0] if _lib_error else "unknown reason" + return ( + f"the Metal splat kernels are not available ({reason}); see " + "torchref.base.electron_density.kernels.mps.compile.last_error()" + ) + + +def mps_kernels_available() -> bool: + """Whether the Metal splat kernels compiled and are ready to dispatch. + + Derived from :func:`why_unavailable` rather than re-testing, so there is one + availability check here, not two that can drift. + """ + return why_unavailable() is None + + +def warmup() -> bool: + """Eagerly trigger compilation (e.g. to move the one-time cost off the + first refinement step). Returns availability.""" + return _get_lib() is not None + + +def clear_cache() -> None: + """Forget the compiled library and failure state (recompiled on next use).""" + global _lib, _lib_failed, _lib_error + _lib = None + _lib_failed = False + _lib_error = None + + +def last_error() -> Optional[Tuple[str, str]]: + """The (message, traceback) of the last compile failure, if any.""" + return _lib_error diff --git a/torchref/base/electron_density/kernels/mps/variable_radius.py b/torchref/base/electron_density/kernels/mps/variable_radius.py new file mode 100644 index 00000000..be5d3f8f --- /dev/null +++ b/torchref/base/electron_density/kernels/mps/variable_radius.py @@ -0,0 +1,202 @@ +"""Metal (MPS) variable-radius electron-density splat, autograd-wrapped. + +``add_isotropic_mps_var`` / ``add_anisotropic_mps_var`` mirror the signatures of +the portable CPU reference (``add_isotropic_plain_var`` / +``add_anisotropic_plain_var`` in ``kernels/cpu/variable_radius.py``) so the +dispatch site in ``main.py`` is a near-copy of the CPU branch. Each delegates to +a ``torch.autograd.Function`` that dispatches the compiled Metal kernels +(one thread per atom) for the forward and analytic backward. + +Gradients flow to ``xyz``, ``adp``/``u``, and ``occ`` (and identity to the input +``density_map``); ``A``/``B`` and the cell matrices receive no gradient -- the +same set as the CUDA kernels. Backward is first-order only (like CUDA); double +backward must use ``force_portable`` (the plain splat). +""" + +from __future__ import annotations + +import torch + +from torchref.base.electron_density.kernels.mps.compile import _get_lib + + +def _r2cut(radius_per_atom): + """Per-atom squared cutoff: the policy radius, used raw. + + This used to round ``radius_per_atom`` up to a whole voxel, to match the old + plain-splat reference whose offset list was built in voxel units. That made the + effective cutoff **grid-dependent** -- inflating it by up to one voxel (~0.4 A, + i.e. ~0.4 sigma at 0.4 A sampling), so Metal truncated systematically later than + the CUDA kernel and ``torchref.sigma_cutoff_ed`` meant different things on the two + devices. The radius arriving from + :mod:`torchref.base.electron_density.radius_policy` is already quantized (0.25 A) + and clamped, so it is used as-is; every backend now shares one cutoff. + """ + return radius_per_atom * radius_per_atom + + +class MetalGridDensity(torch.autograd.Function): + """Isotropic per-atom Metal splat: returns ``density_map + splat``.""" + + @staticmethod + def forward(ctx, density_map, xyz, adp, occ, A, B, r2cut, mask, inv_frac, frac): + lib = _get_lib() + if lib is None: + raise RuntimeError("Metal splat kernels unavailable") + nx, ny, nz = (int(s) for s in density_map.shape) + out = density_map.contiguous().clone() + n = xyz.shape[0] + if n > 0: + lib.iso_splat_fwd( + out.view(-1), + xyz.contiguous(), + adp.contiguous(), + occ.contiguous(), + A.contiguous(), + B.contiguous(), + r2cut.contiguous(), + mask.contiguous(), + inv_frac.contiguous().view(-1), + frac.contiguous().view(-1), + n, nx, ny, nz, + threads=[n], + ) + ctx.save_for_backward(xyz, adp, occ, A, B, r2cut, mask, inv_frac, frac) + ctx.grid_shape = (nx, ny, nz) + return out + + @staticmethod + def backward(ctx, grad_out): + xyz, adp, occ, A, B, r2cut, mask, inv_frac, frac = ctx.saved_tensors + nx, ny, nz = ctx.grid_shape + n = xyz.shape[0] + grad_xyz = torch.zeros_like(xyz) + grad_adp = torch.zeros_like(adp) + grad_occ = torch.zeros_like(occ) + if n > 0: + lib = _get_lib() + lib.iso_splat_bwd( + grad_xyz.view(-1), + grad_adp, + grad_occ, + grad_out.contiguous().view(-1), + xyz.contiguous(), + adp.contiguous(), + occ.contiguous(), + A.contiguous(), + B.contiguous(), + r2cut.contiguous(), + mask.contiguous(), + inv_frac.contiguous().view(-1), + frac.contiguous().view(-1), + n, nx, ny, nz, + threads=[n], + ) + # forward returned density_map + splat -> grad wrt density_map is identity. + # order: density_map, xyz, adp, occ, A, B, r2cut, mask, inv_frac, frac + return (grad_out, grad_xyz, grad_adp, grad_occ, + None, None, None, None, None, None) + + +def add_isotropic_mps_var( + density_map, xyz, adp, occ, A, B, inv_frac_matrix, frac_matrix, radius_per_atom +): + """Isotropic variable-radius Metal splat; adds into ``density_map``. + + Signature is the canonical splat signature, identical to + ``add_isotropic_plain_var`` and ``add_isotropic_cpu_sphere_var`` -- the grid shape + comes from ``density_map`` and the truncation box from the inverse-cell row norms, + so no ``grid_shape_tuple`` or ``voxel_size`` is needed. Both used to be accepted + and ignored. + + The truncation radius is the policy radius, used raw; see :func:`_r2cut`. + """ + r2cut = _r2cut(radius_per_atom) + mask = torch.ones(xyz.shape[0], 5, dtype=xyz.dtype, device=xyz.device) + return MetalGridDensity.apply( + density_map, xyz, adp, occ, A, B, r2cut, mask, inv_frac_matrix, frac_matrix + ) + + +class MetalGridDensityAniso(torch.autograd.Function): + """Anisotropic per-atom Metal splat: returns ``density_map + splat``.""" + + @staticmethod + def forward(ctx, density_map, xyz, u, occ, A, B, r2cut, mask, inv_frac, frac): + lib = _get_lib() + if lib is None: + raise RuntimeError("Metal splat kernels unavailable") + nx, ny, nz = (int(s) for s in density_map.shape) + out = density_map.contiguous().clone() + n = xyz.shape[0] + if n > 0: + lib.aniso_splat_fwd( + out.view(-1), + xyz.contiguous(), + u.contiguous(), + occ.contiguous(), + A.contiguous(), + B.contiguous(), + r2cut.contiguous(), + mask.contiguous(), + inv_frac.contiguous().view(-1), + frac.contiguous().view(-1), + n, nx, ny, nz, + threads=[n], + ) + ctx.save_for_backward(xyz, u, occ, A, B, r2cut, mask, inv_frac, frac) + ctx.grid_shape = (nx, ny, nz) + return out + + @staticmethod + def backward(ctx, grad_out): + xyz, u, occ, A, B, r2cut, mask, inv_frac, frac = ctx.saved_tensors + nx, ny, nz = ctx.grid_shape + n = xyz.shape[0] + grad_xyz = torch.zeros_like(xyz) + grad_u = torch.zeros_like(u) + grad_occ = torch.zeros_like(occ) + if n > 0: + lib = _get_lib() + lib.aniso_splat_bwd( + grad_xyz.view(-1), + grad_u.view(-1), + grad_occ, + grad_out.contiguous().view(-1), + xyz.contiguous(), + u.contiguous(), + occ.contiguous(), + A.contiguous(), + B.contiguous(), + r2cut.contiguous(), + mask.contiguous(), + inv_frac.contiguous().view(-1), + frac.contiguous().view(-1), + n, nx, ny, nz, + threads=[n], + ) + # order: density_map, xyz, u, occ, A, B, r2cut, mask, inv_frac, frac + return (grad_out, grad_xyz, grad_u, grad_occ, + None, None, None, None, None, None) + + +def add_anisotropic_mps_var( + density_map, xyz, u, occ, A, B, inv_frac_matrix, frac_matrix, radius_per_atom +): + """Anisotropic variable-radius Metal splat; adds into ``density_map``. + + Signature is the canonical splat signature, identical to + ``add_anisotropic_plain_var`` and ``add_anisotropic_cpu_sphere_var``. It used to + take ``real_space_grid`` first and ``voxel_size`` last; neither was read -- the + grid shape comes from ``density_map``, which is what gets forwarded to + ``MetalGridDensityAniso``. The old docstring claimed ``real_space_grid`` was "used + only for its grid shape", which had stopped being true. + + Each atom is truncated at its per-axis bounding box with a sphere cull at the + per-atom radius -- the same canonical cutoff the CUDA and fused-CPU kernels apply. + """ + r2cut = _r2cut(radius_per_atom) + mask = torch.ones(xyz.shape[0], 5, dtype=xyz.dtype, device=xyz.device) + return MetalGridDensityAniso.apply( + density_map, xyz, u, occ, A, B, r2cut, mask, inv_frac_matrix, frac_matrix + ) diff --git a/torchref/base/electron_density/main.py b/torchref/base/electron_density/main.py index 621ad98e..c47ba21d 100644 --- a/torchref/base/electron_density/main.py +++ b/torchref/base/electron_density/main.py @@ -1,36 +1,46 @@ """ -Central electron density building, dispatched solely by the shared ``Engine``. +Central electron density building, dispatched from one declarative table. +One truncation contract, every backend +------------------------------------- Every atom is splatted at its own per-atom truncation radius (``N_sigma * sigma_eff``, with ``N_sigma = torchref.sigma_cutoff_ed``); there is -no single global splat radius. The capability-based ``Engine`` in -:mod:`torchref.utils.triton_dispatch` (AUTO/TRITON/EAGER) is the *only* switch -selecting which variable-radius kernel runs — there is no environment-variable -dispatch and no parallel "tier" knobs: - -- ``Engine.AUTO`` — fastest available per device, all variable-radius: - CUDA+float32 -> the work-queue Triton kernels - (``WorkQueueGridDensity`` / ``WorkQueueGridDensityAniso``); CPU -> the - grouped-separable variable-radius splat - (``add_isotropic_cpu_separable_var`` / ``add_anisotropic_cpu_var``); - everything else (CUDA+float64, MPS) -> the portable plain-scatter - variable-radius splat (``add_isotropic_plain_var`` / - ``add_anisotropic_plain_var``). On a Triton kernel failure under AUTO it falls - through to the plain-scatter variable-radius splat. -- ``Engine.EAGER`` — the portable plain-scatter variable-radius splat on every - device. Double-differentiable; use it for Hessians / debugging. Force it with - ``with use_engine(Engine.EAGER): ...``. -- ``Engine.TRITON`` — force the CUDA work-queue Triton kernel (raises if not - CUDA+float32). - -The production variable-radius splats live in -``kernels/cuda/variable_radius.py`` and ``kernels/cpu/variable_radius.py``; the -per-atom radius policy is in :mod:`torchref.base.electron_density.radius_policy`. -Several legacy fixed-radius kernels are re-imported here so the historical -``torchref.base.electron_density.main`` namespace is unchanged and they remain -callable directly for benchmarking, but they are *not* on the production -dispatch path. This module keeps only the ``Engine``-based variable-radius -dispatch. +no single global splat radius. Crucially, every production kernel now applies the +*same* cutoff, so ``sigma_cutoff_ed`` means the same thing on every device: + + voxel v receives atom i's density iff ``||w||^2 <= r_i^2``, where ``w`` is the + minimum-image **Cartesian atom->voxel** vector (the sphere is centred on the + atom, not on its nearest grid node) and ``r_i`` is the raw + :mod:`~torchref.base.electron_density.radius_policy` radius, + +enumerated over the triclinic-correct per-axis box +``ceil(r_i * n_axis * ||inv_frac row_axis||)``. No grid-dependent requantization of +the radius, no diagonal-metric approximation, no cube. Historically the backends +disagreed here -- Metal inflated the radius to a whole voxel, the portable CPU +splat used a node-centred diagonal metric, and the CPU fast path splatted a cube -- +by amounts comparable to or larger than the truncation error itself. + +Backend dispatch +---------------- +Which kernel runs is decided from a table, not from an if/elif ladder: see +:data:`torchref.base.electron_density._backends.DENSITY_BACKENDS`. That table is the only +place the criteria are written down -- device, dtype, how availability is probed, and whether +a runtime failure may degrade -- so it is also the only place to look, or to edit when adding +a backend. The mechanics of reading it live in :mod:`torchref.utils.backends`. + +Selection needs no configuration: the fastest kernel that can run for the given device and +dtype is chosen, and an accelerator that is missing or throws degrades to the portable splat +(with a warning). There is no environment-variable dispatch and no "tier" knobs. + +The one override is ``force_portable``, which pins the portable reference splat. Pass it per +call or scope it with ``with use_portable(): ...``. Its purpose is the single failure that +automatic fallback cannot detect -- an accelerator kernel that runs and returns *wrong +numbers* rather than raising. + +The production splats live in ``kernels/cuda/variable_radius.py``, +``kernels/cpu/sphere_splat.py``, ``kernels/mps/variable_radius.py`` and (portable) +``kernels/cpu/variable_radius.py``; the per-atom radius policy is in +:mod:`torchref.base.electron_density.radius_policy`. """ from typing import Optional @@ -38,38 +48,17 @@ import torch from torchref.config import get_float_dtype, get_sigma_cutoff_ed -from torchref.utils.triton_dispatch import Engine, get_engine, should_use_triton +from torchref.utils.backends import run_or_degrade, select -# --- Shared splat helpers (re-imported to preserve this namespace; reused by the -# variable-radius kernels and, for _get_radius_offsets, by scaling/solvent.py) --- -from torchref.base.electron_density.kernels.offsets import _get_radius_offsets -from torchref.base.electron_density.kernels.cpu.scatter_dispatch import ( - _do_structured_scatter, - _get_cpp_scatter, -) -from torchref.base.electron_density.kernels.cpu.separable import _separable_density -from torchref.base.electron_density.kernels.cpu.aniso import _aniso_density_cube +from torchref.base.electron_density._backends import DENSITY_BACKENDS -# --- Per-atom variable-radius density path --- -# The splat radius is no longer a single scalar; each atom is truncated at its own -# N_sigma * sigma_eff radius (N_sigma = torchref.sigma_cutoff_ed). CUDA+float32 uses -# the variable-radius Triton kernels (WorkQueueGridDensity{,Aniso}); CPU + AUTO uses -# the grouped-separable structured-scatter splat; everything else (EAGER any device, -# CUDA float64, MPS) uses the portable plain-scatter splat. +# Re-imported to preserve this namespace: ``scaling/solvent.py`` imports +# ``_get_radius_offsets`` from here, not from its defining module. +from torchref.base.electron_density.kernels.offsets import _get_radius_offsets from torchref.base.electron_density.radius_policy import ( per_atom_radius_aniso, per_atom_radius_iso, ) -from torchref.base.electron_density.kernels.cuda.variable_radius import ( - WorkQueueGridDensity, - WorkQueueGridDensityAniso, -) -from torchref.base.electron_density.kernels.cpu.variable_radius import ( - add_anisotropic_cpu_var, - add_anisotropic_plain_var, - add_isotropic_cpu_separable_var, - add_isotropic_plain_var, -) def build_electron_density( @@ -88,19 +77,14 @@ def build_electron_density( A_aniso: Optional[torch.Tensor] = None, B_aniso: Optional[torch.Tensor] = None, dtype: torch.dtype = None, + force_portable: Optional[bool] = None, ) -> torch.Tensor: """ Build an electron density map from atomic parameters. - Dispatches the isotropic and anisotropic splats through the shared - ``Engine`` (see the module docstring). Each atom is splatted at its own - per-atom truncation radius (``N_sigma * sigma_eff``, with - ``N_sigma = torchref.sigma_cutoff_ed``) via the per-atom variable-radius - kernels: the CUDA float32 work-queue kernels - (``WorkQueueGridDensity{,Aniso}``) under ``Engine.AUTO``/``Engine.TRITON``, - the grouped-separable variable-radius splat on CPU+AUTO, and the portable - plain-scatter variable-radius splat everywhere else (``Engine.EAGER``, CUDA - float64, MPS). + Each atom is splatted at its own per-atom truncation radius + (``N_sigma * sigma_eff``, with ``N_sigma = torchref.sigma_cutoff_ed``). Which kernel + does the splatting is decided from ``DENSITY_BACKENDS``; see the module docstring. Parameters ---------- @@ -119,8 +103,11 @@ def build_electron_density( frac_matrix : torch.Tensor Fractional-to-Cartesian matrix, shape (3, 3). voxel_size : torch.Tensor - Voxel dimensions, shape (3,). The per-atom truncation radius is derived - internally from each atom's B/U and ``torchref.sigma_cutoff_ed``. + Voxel dimensions, shape (3,). **Unused by every splat** -- the per-atom + truncation radius comes from each atom's B/U and ``torchref.sigma_cutoff_ed`` + via :mod:`~torchref.base.electron_density.radius_policy`, and the enumeration + box from ``inv_frac_matrix``. Retained because ``SfFFT`` passes it + positionally; dropping it is a wider API change. xyz_aniso : torch.Tensor, optional Anisotropic atom positions, shape (n_aniso, 3). u_aniso : torch.Tensor, optional @@ -132,6 +119,11 @@ def build_electron_density( dtype : torch.dtype, optional Float dtype for the density map. Defaults to the configured float dtype (``get_float_dtype()``), which may be float64. + force_portable : bool, optional + Pin the portable reference splat instead of the fastest available kernel. ``None`` + defers to the process-wide setting (see ``torchref.utils.use_portable``). Use it to + check whether an accelerator kernel is producing wrong numbers -- the one failure + automatic fallback cannot detect. Returns ------- @@ -150,7 +142,6 @@ def build_electron_density( # --- isotropic atoms --- if len(xyz_iso) > 0: density_map = _add_isotropic( - real_space_grid, density_map, xyz_iso, adp_iso, @@ -159,13 +150,12 @@ def build_electron_density( B_iso, inv_frac_matrix, frac_matrix, - voxel_size, + force_portable=force_portable, ) # --- anisotropic atoms --- if xyz_aniso is not None and len(xyz_aniso) > 0: density_map = _add_anisotropic( - real_space_grid, density_map, xyz_aniso, u_aniso, @@ -174,7 +164,7 @@ def build_electron_density( B_aniso, inv_frac_matrix, frac_matrix, - voxel_size, + force_portable=force_portable, ) return density_map @@ -186,7 +176,6 @@ def build_electron_density( def _add_isotropic( - real_space_grid, density_map, xyz, adp, @@ -195,70 +184,28 @@ def _add_isotropic( B, inv_frac_matrix, frac_matrix, - voxel_size, + force_portable=None, ): - """Add isotropic atoms with a per-atom variable radius. ``Engine`` is the switch. + """Add isotropic atoms with a per-atom variable radius. The per-atom radius is ``clamp(ceil(N_sigma * sigma_eff), [2,7])`` with ``N_sigma = torchref.sigma_cutoff_ed``. - - ``should_use_triton`` (CUDA + float32 + Triton, engine AUTO/TRITON) -> the - variable-radius Triton kernel (``WorkQueueGridDensity``). On kernel failure - under AUTO it falls through to the portable splat; under ``Engine.TRITON`` - it raises (never silently degrade). - - CPU + AUTO -> the fast C++-scatter grouped-separable splat. - - Everything else (``Engine.EAGER`` on any device, CUDA float64, MPS) -> the - portable plain-``scatter_add`` grouped splat: identical per-atom radius, - double-differentiable, float64-capable, device-agnostic. + Which kernel runs, and what happens if it fails, are read from ``DENSITY_BACKENDS`` + rather than restated here -- there is no branch in this function to keep in sync with + the table. Every path applies the identical spherical cutoff, so the choice affects + speed, not result. - Every path truncates each atom at its own ``N_sigma * sigma_eff`` radius. + Only the first six arguments carry the device/dtype contract; the table names them. """ - n_sigma = get_sigma_cutoff_ed() - radius_per_atom = per_atom_radius_iso(adp, B, n_sigma=n_sigma) - - if should_use_triton(xyz): - try: - r2cut = radius_per_atom * radius_per_atom - coeff_mask = torch.ones(xyz.shape[0], 5, dtype=xyz.dtype, device=xyz.device) - # kernel accumulates into (a copy of) density_map -> no extra grid buffer + add - return WorkQueueGridDensity.apply( - density_map, - real_space_grid, - xyz, - adp, - occ, - A, - B, - r2cut, - coeff_mask, - inv_frac_matrix, - frac_matrix, - ) - except Exception: - if get_engine() is Engine.TRITON: - raise - # AUTO: fall through to the portable splat - - grid_shape_tuple = real_space_grid.shape[:3] - # The C++-scatter fast path is float32-only; non-float32 (float64 config, - # complex-derived double grids) falls through to the portable plain splat. - if ( - get_engine() is Engine.AUTO - and density_map.device.type == "cpu" - and density_map.dtype == torch.float32 - ): - return add_isotropic_cpu_separable_var( - density_map, xyz, adp, occ, A, B, - inv_frac_matrix, frac_matrix, grid_shape_tuple, voxel_size, radius_per_atom, - ) - return add_isotropic_plain_var( - density_map, xyz, adp, occ, A, B, - inv_frac_matrix, frac_matrix, grid_shape_tuple, voxel_size, radius_per_atom, - ) + radius_per_atom = per_atom_radius_iso(adp, B, n_sigma=get_sigma_cutoff_ed()) + args = (density_map, xyz, adp, occ, A, B, + inv_frac_matrix, frac_matrix, radius_per_atom) + backend = select(DENSITY_BACKENDS, args, force_portable=force_portable) + return run_or_degrade(DENSITY_BACKENDS, backend, False, *args) def _add_anisotropic( - real_space_grid, density_map, xyz, u, @@ -267,58 +214,20 @@ def _add_anisotropic( B, inv_frac_matrix, frac_matrix, - voxel_size, + force_portable=None, ): """Add anisotropic atoms with a per-atom variable radius (mirrors the iso path). - The per-atom radius is the isotropic bounding radius of the ellipsoid - (largest principal axis, ``per_atom_radius_aniso``). + The per-atom radius is the isotropic bounding radius of the ellipsoid (largest + principal axis, ``per_atom_radius_aniso``). Every path culls on the Euclidean sphere at + that radius and evaluates the Mahalanobis form inside it -- one contract, as for the + isotropic pass. - CUDA+float32 (engine permitting) -> ``WorkQueueGridDensityAniso``. CPU + AUTO - -> the fast C++-scatter grouped box-splat ``add_anisotropic_cpu_var``. - Everything else (``Engine.EAGER`` on any device, CUDA float64, MPS) -> the - portable plain-``scatter_add`` box-splat ``add_anisotropic_plain_var`` - (double-diff, float64-capable, device-agnostic). All paths use the per-atom - radius (isotropic bounding sphere of the ellipsoid). + Same table, same selection: the two passes differ only in the radius policy and in + which variant of each kernel the table names. """ - n_sigma = get_sigma_cutoff_ed() - radius_per_atom = per_atom_radius_aniso(B, u, n_sigma=n_sigma) - - if should_use_triton(xyz): - try: - r2cut = radius_per_atom * radius_per_atom - coeff_mask = torch.ones(xyz.shape[0], 5, dtype=xyz.dtype, device=xyz.device) - # kernel accumulates into (a copy of) density_map -> no extra grid buffer + add - return WorkQueueGridDensityAniso.apply( - density_map, - real_space_grid, - xyz, - u, - occ, - A, - B, - r2cut, - coeff_mask, - inv_frac_matrix, - frac_matrix, - ) - except Exception: - if get_engine() is Engine.TRITON: - raise - # AUTO: fall back to the portable splat - - # The C++-scatter fast path is float32-only; non-float32 falls through to - # the portable plain splat (see _add_isotropic). - if ( - get_engine() is Engine.AUTO - and density_map.device.type == "cpu" - and density_map.dtype == torch.float32 - ): - return add_anisotropic_cpu_var( - real_space_grid, density_map, xyz, u, occ, A, B, - inv_frac_matrix, frac_matrix, radius_per_atom, voxel_size, - ) - return add_anisotropic_plain_var( - real_space_grid, density_map, xyz, u, occ, A, B, - inv_frac_matrix, frac_matrix, radius_per_atom, voxel_size, - ) + radius_per_atom = per_atom_radius_aniso(B, u, n_sigma=get_sigma_cutoff_ed()) + args = (density_map, xyz, u, occ, A, B, + inv_frac_matrix, frac_matrix, radius_per_atom) + backend = select(DENSITY_BACKENDS, args, force_portable=force_portable) + return run_or_degrade(DENSITY_BACKENDS, backend, True, *args) diff --git a/torchref/base/electron_density/radius_policy.py b/torchref/base/electron_density/radius_policy.py index ffcb5af4..0b280559 100644 --- a/torchref/base/electron_density/radius_policy.py +++ b/torchref/base/electron_density/radius_policy.py @@ -48,6 +48,49 @@ def _u6_to_u3(u: torch.Tensor) -> torch.Tensor: return U3 +def _max_eig_sym3(A: torch.Tensor) -> torch.Tensor: + """Largest eigenvalue of a batch of symmetric 3x3 matrices, in closed form. + + Uses the analytic (trigonometric) solution of the characteristic cubic + (Smith 1961) so no ``torch.linalg.eigvalsh`` is needed -- that op is + unimplemented on MPS, and the eigendecomposition here feeds only the + (quantized) splat radius, so an exact decomposition is overkill. + + Parameters + ---------- + A : torch.Tensor + Symmetric matrices, shape (n, 3, 3). + + Returns + ------- + torch.Tensor + Largest eigenvalue per matrix, shape (n,). + """ + a00 = A[:, 0, 0]; a11 = A[:, 1, 1]; a22 = A[:, 2, 2] + a01 = A[:, 0, 1]; a02 = A[:, 0, 2]; a12 = A[:, 1, 2] + + p1 = a01 * a01 + a02 * a02 + a12 * a12 + q = (a00 + a11 + a22) / 3.0 + p2 = (a00 - q) ** 2 + (a11 - q) ** 2 + (a22 - q) ** 2 + 2.0 * p1 + p = torch.sqrt((p2 / 6.0).clamp(min=1e-30)) + + # B = (A - q I) / p ; r = det(B) / 2 in [-1, 1] + b00 = (a00 - q) / p; b11 = (a11 - q) / p; b22 = (a22 - q) / p + b01 = a01 / p; b02 = a02 / p; b12 = a12 / p + detB = ( + b00 * (b11 * b22 - b12 * b12) + - b01 * (b01 * b22 - b12 * b02) + + b02 * (b01 * b12 - b11 * b02) + ) + r = (detB / 2.0).clamp(-1.0, 1.0) + phi = torch.acos(r) / 3.0 + eig_max = q + 2.0 * p * torch.cos(phi) + + # Diagonal matrices (p1 == 0): eigenvalues are the diagonal entries. + diag_max = torch.maximum(torch.maximum(a00, a11), a22) + return torch.where(p1 <= 1e-30, diag_max, eig_max) + + def sigma_eff_iso(adp: torch.Tensor, B_widths: torch.Tensor) -> torch.Tensor: """``sigma_eff_i = sqrt((max_k B_widths[i,k] + adp_i) / 8pi^2)``, shape (n,). @@ -93,7 +136,8 @@ def sigma_eff_aniso(B_widths: torch.Tensor, u: torch.Tensor) -> torch.Tensor: Anisotropic U parameters [U11,U22,U33,U12,U13,U23], shape (n, 6). """ b_form = B_widths.detach().max(dim=1).values # broadest ITC92 width - lam_max = torch.linalg.eigvalsh(_u6_to_u3(u.detach())).max(dim=1).values + # Closed-form largest eigenvalue (no eigvalsh -> runs natively on MPS). + lam_max = _max_eig_sym3(_u6_to_u3(u.detach())) return torch.sqrt((b_form / EIGHT_PI2 + lam_max).clamp(min=1e-6)) diff --git a/torchref/base/kernels/__init__.py b/torchref/base/kernels/__init__.py index 1ab60db8..74c42528 100644 --- a/torchref/base/kernels/__init__.py +++ b/torchref/base/kernels/__init__.py @@ -16,16 +16,8 @@ warmup, get_cache_dir, clear_cache, - _HAS_TRITON, ) -try: - from torchref.base.electron_density.kernels import ( # noqa: F401 - fused_add_to_map_gpu, - ) -except ImportError: - pass - __all__ = [ "vectorized_add_to_map", "build_electron_density", @@ -34,5 +26,17 @@ "warmup", "get_cache_dir", "clear_cache", - "fused_add_to_map_gpu", ] + +# Optional, and only advertised when it actually resolved -- listing it +# unconditionally made ``from torchref.base.kernels import *`` raise on a Triton-less +# host. The former ``_HAS_TRITON`` re-export is gone: nothing read it, and +# ``torchref.utils.triton_available()`` is the answer to that question. +try: + from torchref.base.electron_density.kernels import ( # noqa: F401 + fused_add_to_map_gpu, + ) + + __all__.append("fused_add_to_map_gpu") +except ImportError: + pass diff --git a/torchref/base/reciprocal/symmetry.py b/torchref/base/reciprocal/symmetry.py index e99e0f03..902afdbe 100644 --- a/torchref/base/reciprocal/symmetry.py +++ b/torchref/base/reciprocal/symmetry.py @@ -34,7 +34,7 @@ import numpy as np import torch -from torchref.config import get_float_dtype +from torchref.config import canonical_device, get_float_dtype from torchref.utils.autograd_ops import gather_with_index_add from .grid_operations import extract_structure_factor_from_grid @@ -251,7 +251,12 @@ def __init__( grid_shape: tuple, device: Optional[torch.device] = None, ): - self.device = device or hkl.device + # ``is not None``, not ``or``: ``device=0`` means cuda:0/mps:0 and is + # falsy, so ``or`` silently discarded it. ``hkl`` is a bare tensor, so + # this reads its device rather than going through ``resolve_device``. + self.device = canonical_device( + device if device is not None else hkl.device + ) self.hkl = hkl.to(device=self.device) self.symmetry = symmetry self.n_ops = symmetry.n_ops diff --git a/torchref/base/targets/_dispatch.py b/torchref/base/targets/_dispatch.py index b34b14bb..4a925eb4 100644 --- a/torchref/base/targets/_dispatch.py +++ b/torchref/base/targets/_dispatch.py @@ -1,31 +1,99 @@ -"""Shared dispatch helper for the target math functions. +"""Shared dispatch gate for the target math functions. -When called on a CUDA float32 tensor and Triton is importable, the math -functions in this package transparently route to their Triton-kernel -implementations in :mod:`torchref.base.targets.triton`. CPU tensors, -non-float32 tensors, or environments without Triton fall back to the -plain eager implementation. +When called on a CUDA float32 tensor and the Triton kernels are importable, the math +functions in this package route to their implementations in +:mod:`torchref.base.targets.triton`. CPU tensors, non-float32 tensors, and environments +without a usable Triton fall back to the plain eager implementation. -Selection is governed by the shared, capability-based ``Engine`` in -:mod:`torchref.utils.triton_dispatch` (no environment variables). To force -the eager path for an A/B comparison or to sidestep a flaky Triton install:: +The criteria live in :data:`TARGET_BACKENDS` rather than in a predicate body, for the same +reason as the density and direct-summation tables: device, dtype and availability stated once, +as data, next to the kernels they select. - from torchref.utils import use_engine, Engine - with use_engine(Engine.EAGER): +This is a **gate-only** table -- the twelve call sites each do their own +``from .triton. import `` three lines from the ``if``, which is more legible than a +registry lookup for a single-kernel choice. The table's job is the decision, not the +dispatch, so its rows carry no ``kernel``. + +To force the eager path for an A/B comparison or to sidestep a flaky Triton install:: + + from torchref.utils import use_portable + with use_portable(): ... """ +from typing import Optional + import torch -from torchref.utils.triton_dispatch import should_use_triton +from torchref.utils.backends import Backend, BackendTable, triton_available, will_use + +_THIS = "torchref.base.targets._dispatch" + + +def why_unavailable() -> Optional[str]: + """``None`` if the target Triton kernels can run, else why they cannot. + + Closes a real gap. ``triton_available()`` answers "is the ``triton`` package + importable", which is not the same question as "do *these* kernels import" -- and + nothing asked the second one. Each of the twelve call sites does an unguarded + function-local import, so a Triton present but skewed against the installed driver or + LLVM raised straight through a refinement step, where the density and direct-summation + paths would have degraded. Importing the package is the shared prerequisite for all + twelve, so one probe covers them. + + Deliberately does *not* re-check ``torch.cuda.is_available()``. Selection is two-phase, + so the device criterion has already matched by the time a probe runs -- and the presence + of a CUDA tensor is stronger evidence than the query. Asking again would also make this + probe disagree with ``triton_available()`` on a host that has Triton installed but no + GPU, which is an ordinary CI configuration. + """ + if not triton_available(): + return "triton is not importable" + try: + import torchref.base.targets.triton # noqa: F401 + except Exception as exc: # noqa: BLE001 - an import failure is a reason, not a crash + return ( + "the target Triton kernels could not be imported " + f"({type(exc).__name__}: {exc})" + ) + return None + + +TARGET_BACKENDS = BackendTable( + name="target math", + backends=( + Backend( + name="triton", + kernel=None, # gate-only; see the module docstring + device="cuda", + dtypes=(torch.float32,), + probe=(_THIS, "why_unavailable"), + expect_available="cuda", + # Availability is handled by the probe above, so this governs only a kernel + # that imported and then threw -- a bug in pure math on already-validated + # tensors. Degrading would buy a silent ~50x slowdown with subtly different + # numbers, which is worse than the exception. + on_failure="raise", + second_order=False, + ), + Backend( + name="eager", + kernel=None, + # METAL is here because there are no Metal target kernels: at this site it has + # to mean "run eager". TRITON is absent, which is what makes it strict. + expect_available="always", + on_failure="raise", + second_order=True, + ), + ), +) def use_triton(*tensors: torch.Tensor) -> bool: """Decide whether to route a call to the Triton kernel. - Thin wrapper over :func:`torchref.utils.triton_dispatch.should_use_triton` - using the process-wide engine: Triton is used only when the engine permits - it and every probed tensor is CUDA float32 (the only configuration the - kernels are written for). + Asks the ``triton`` row of :data:`TARGET_BACKENDS` whether it is the backend that would + actually be selected. ``None`` entries among ``tensors`` are ignored, so a caller may pass + optional inputs straight through. """ - return should_use_triton(*tensors) + return will_use(TARGET_BACKENDS, "triton", tensors) diff --git a/torchref/base/targets/xray_ml_sigmaa.py b/torchref/base/targets/xray_ml_sigmaa.py index cb92f092..560f9c1a 100644 --- a/torchref/base/targets/xray_ml_sigmaa.py +++ b/torchref/base/targets/xray_ml_sigmaa.py @@ -36,6 +36,8 @@ import numpy as np import torch +from torchref.config import get_float_dtype + # ===================================================================== # Loss math (mean = |Fc|, variance = epsilon*beta) # ===================================================================== @@ -110,21 +112,34 @@ def epsilon_from_hkl(hkl: torch.Tensor, spacegroup) -> torch.Tensor: Mirrors ``ReciprocalSymmetry.get_epsilon`` (Friedel-aware) but works directly on the scattered HKL list. Returns ones if ``spacegroup`` is None or lacks ``apply_to_hkl``. + + Always returns on ``hkl.device``, whatever device the space group's symmetry + matrices live on: the caller multiplies this against per-reflection data + sitting beside ``hkl``. """ n = hkl.shape[0] - ones = torch.ones(n, device=hkl.device) + float_dtype = get_float_dtype() if spacegroup is None or not hasattr(spacegroup, "apply_to_hkl"): - return ones - try: - with torch.no_grad(): - Hs = spacegroup.apply_to_hkl(hkl.to(torch.float64)) # (N,3,ops) - h0 = hkl.to(torch.float64).unsqueeze(-1) # (N,3,1) - same = (Hs == h0).all(dim=1) - friedel = (Hs == -h0).all(dim=1) - eps = (same | friedel).sum(dim=1).clamp(min=1).to(torch.get_default_dtype()) - return eps.to(hkl.device) - except Exception: - return ones + return torch.ones(n, device=hkl.device, dtype=float_dtype) + + with torch.no_grad(): + # Configured float dtype, not float64: MPS has no float64 and casting + # there raises. Symmetry arithmetic on Miller indices is exact in + # float32 (integer-valued rotation matrices, small indices), so the + # exact `==` comparisons below remain valid. + # + # ``apply_to_hkl`` moves its input onto the matrices' device, so build + # ``h`` there too -- otherwise ``Hs`` and ``h0`` land on different + # devices and the comparisons below raise. The space group wins for the + # arithmetic; the result is handed back on the caller's device. + sym_device = getattr(spacegroup, "matrices", hkl).device + h = hkl.to(device=sym_device, dtype=float_dtype) + Hs = spacegroup.apply_to_hkl(h) # (N,3,ops) + h0 = h.unsqueeze(-1) # (N,3,1) + same = (Hs == h0).all(dim=1) + friedel = (Hs == -h0).all(dim=1) + eps = (same | friedel).sum(dim=1).clamp(min=1).to(float_dtype) + return eps.to(hkl.device) # ===================================================================== diff --git a/torchref/cli/_common.py b/torchref/cli/_common.py index f91b8542..67a90602 100644 --- a/torchref/cli/_common.py +++ b/torchref/cli/_common.py @@ -633,10 +633,9 @@ def load_model( ModelFT """ from torchref import ModelFT - from torchref.config import get_default_device + from torchref.config import normalize_device - if device is None: - device = get_default_device() + device = normalize_device(device) model = ModelFT(max_res=max_res, device=device, verbose=verbose) suffix = Path(path).suffix.lower() if suffix in (".cif", ".mmcif"): @@ -674,10 +673,9 @@ def load_reflection_data( ReflectionData """ from torchref import ReflectionData - from torchref.config import get_default_device + from torchref.config import normalize_device - if device is None: - device = get_default_device() + device = normalize_device(device) data = ReflectionData(device=str(device), verbose=verbose) suffix = Path(path).suffix.lower() if suffix in (".cif",): diff --git a/torchref/cli/validate_ded.py b/torchref/cli/validate_ded.py index 60156f6a..1f457df6 100644 --- a/torchref/cli/validate_ded.py +++ b/torchref/cli/validate_ded.py @@ -292,10 +292,9 @@ def setup_ded_context( from torchref.symmetry.grid_utils import calculate_optimal_grid_size from torchref.symmetry.reciprocal_symmetry import expand_hkl - if device is None: - from torchref.config import get_default_device + from torchref.config import normalize_device - device = get_default_device() + device = normalize_device(device) # Load and scale data_dark = load_reflection_data( diff --git a/torchref/config.py b/torchref/config.py index 94b253f2..33eceb62 100644 --- a/torchref/config.py +++ b/torchref/config.py @@ -358,6 +358,65 @@ def _auto_detect_device() -> torch.device: return torch.device("cpu") +def canonical_device(dev): + """Return ``dev`` with its default index filled in. + + ``torch.device('cuda') != torch.device('cuda:0')`` even though both name the + same physical device. Devices read back off a real tensor always carry an + index, so any comparison between a *requested* device and an *observed* one + needs a shared normal form -- otherwise ``obj.device == tensor.device`` is + False on a freshly constructed object and True after its first ``.to()``. + + ``cpu`` deliberately stays bare: ``torch.empty(0, device='cpu').device`` has + no index, so canonicalising it to ``cpu:0`` would recreate the very mismatch + this function exists to remove. + + This is the allocation-free form of ``torch.empty(0, device=d).device``, + which matters because it is called on constructor and comparison paths. + Deliberately *not* memoised: ``torch.cuda.current_device()`` is mutable via + ``torch.cuda.set_device``, so a cache would freeze a stale answer. + + Parameters + ---------- + dev : torch.device or str or int or None + Device to normalise. ``None`` passes through so callers can chain. + + Returns + ------- + torch.device or None + """ + if dev is None: + return None + d = dev if isinstance(dev, torch.device) else torch.device(dev) + if d.index is not None: + return d + if d.type == "cpu": + return d + if d.type == "cuda": + return torch.device("cuda", torch.cuda.current_device()) + if d.type == "mps": + return torch.device("mps", 0) + # Unknown/future backend: fall back to asking PyTorch directly. + return torch.empty(0, device=d).device + + +def normalize_device(dev=None) -> torch.device: + """Coerce a user-supplied ``device`` argument to a canonical device. + + ``None`` resolves to :func:`get_default_device`. This is the pure, + side-effect-free counterpart of :func:`torchref.utils.resolve_device`: + + * one device source (or none) -> ``normalize_device`` + * several device-bearing inputs to reconcile -> ``resolve_device`` + + ``resolve_device`` *moves* the objects it is given, so using it for a + single-input constructor with nothing to reconcile is overreach. + """ + if dev is None: + return get_default_device() + return canonical_device(dev) + + class DeviceConfig: """ Device configuration with property-based access. @@ -381,7 +440,10 @@ class DeviceConfig: def __init__(self): override = os.environ.get("TORCHREF_DEVICE", "auto").lower() if override == "auto": - self._device = _auto_detect_device() + # Canonicalise here too, not just in ``_coerce``: the ``auto`` + # branch bypasses ``_coerce`` entirely, and it is the branch almost + # every user takes. + self._device = canonical_device(_auto_detect_device()) else: self._device = self._coerce(override) self._warn_if_mps_dtype_mismatch() @@ -408,7 +470,7 @@ def _coerce(value) -> torch.device: hasattr(torch.backends, "mps") and torch.backends.mps.is_available() ): raise RuntimeError("MPS requested but not available on this system.") - return dev + return canonical_device(dev) def _warn_if_mps_dtype_mismatch(self) -> None: if self._device.type == "mps" and dtypes.float == torch.float64: diff --git a/torchref/io/datasets/base.py b/torchref/io/datasets/base.py index dd7013b3..53271a74 100644 --- a/torchref/io/datasets/base.py +++ b/torchref/io/datasets/base.py @@ -20,7 +20,7 @@ import gemmi import torch -from torchref.config import get_default_device +from torchref.config import get_default_device, normalize_device from torchref.symmetry import Cell from torchref.utils.device_mixin import DeviceMovementMixin @@ -121,9 +121,10 @@ class CrystalDataset(DeviceMovementMixin): def __post_init__(self): """Initialize non-field attributes after dataclass init.""" - # Ensure device is a torch.device object - if isinstance(self.device, str): - object.__setattr__(self, "device", torch.device(self.device)) + # Canonicalise, don't merely coerce: ``torch.device("mps")`` has no + # index and so compares unequal to the ``mps:0`` every tensor reports, + # which makes ``resolve_device``'s equality checks misfire. + object.__setattr__(self, "device", normalize_device(self.device)) # Import here to avoid circular imports from torchref.utils.utils import TensorMasks @@ -204,8 +205,7 @@ def _from_state(cls, state: Dict[str, Any], device=None) -> "CrystalDataset": """ from torchref.utils.utils import TensorMasks - if device is None: - device = get_default_device() + device = normalize_device(device) # Extract masks before creating object masks_state = state.pop("masks", {}) diff --git a/torchref/io/datasets/fcalc_data.py b/torchref/io/datasets/fcalc_data.py index d5139792..2184b4e1 100644 --- a/torchref/io/datasets/fcalc_data.py +++ b/torchref/io/datasets/fcalc_data.py @@ -14,7 +14,7 @@ import pandas as pd import torch -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_default_device, get_float_dtype, normalize_device from torchref.symmetry import Cell, SpaceGroup, SpaceGroupLike from .base import CrystalDataset @@ -131,13 +131,21 @@ def from_cell_and_resolution( from torchref.base.reciprocal import get_d_spacing - if device is None: - device = get_default_device() + # ``Cell`` is the unit-cell carrier: when one is handed in and no device + # is requested, follow it rather than the global default -- otherwise a + # CPU-default host silently relocates the caller's cell. + if device is None and isinstance(cell, Cell): + device = cell.device + device = normalize_device(device) if dtype is None: dtype = get_float_dtype() # Handle Cell input - convert to Cell object if needed if isinstance(cell, Cell): + # ``Cell.to`` is in-place, so this moves the *caller's* object when + # an explicit device disagrees with it. That is the documented + # resolve_device contract; the branch above avoids it entirely when + # no device was requested. cell_obj = cell.to(device=device) cell_tensor = cell_obj.data else: @@ -376,8 +384,7 @@ def _from_state(cls, state: Dict[str, Any], device=None) -> "FcalcDataset": """ from torchref.utils.utils import TensorMasks - if device is None: - device = get_default_device() + device = normalize_device(device) # Extract masks before creating object masks_state = state.pop("masks", {}) diff --git a/torchref/io/datasets/reflection_data.py b/torchref/io/datasets/reflection_data.py index 4fb74f60..1daebeba 100644 --- a/torchref/io/datasets/reflection_data.py +++ b/torchref/io/datasets/reflection_data.py @@ -8,7 +8,8 @@ import warnings from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union import numpy as np import pandas as pd @@ -17,7 +18,7 @@ from torchref.base import math_torch from torchref.base.french_wilson import FrenchWilson -from torchref.config import dtypes, get_default_device +from torchref.config import dtypes, get_default_device, normalize_device from torchref.io import cif, mtz from torchref.io.datasets.base import CrystalDataset from torchref.symmetry import Cell, SpaceGroup @@ -331,6 +332,162 @@ def _tv(t): self._corrected_fp = fp return self._corrected_cache + # ===================== per-reflection field reindexing ===================== + # + # Any routine that changes the HKL set (expand onto a reference grid, remap, + # symmetry reduction) must carry EVERY per-reflection field along, not a + # hand-maintained subset. Historically ``validate_hkl`` / ``remap`` / + # ``reduce_to_spacegroup`` each hardcoded their own field list and silently + # dropped fields added later (notably ``hkl_anomalous``, read by + # ``hkl_for_sf``), which broke difference refinement on mismatched reflection + # files. The single source of truth below is shared by all of them. + + # Fill value used for MISSING output rows (index == -1) per field. Missing + # rows are always masked out downstream (``hkl_present`` / ``missing``), so + # these fills only need to be shape-correct and non-poisonous. Fields not + # listed default to 0. NOTE: ``hkl_anomalous`` is special-cased (filled with + # the reference HKL row, never 0 which would be a spurious Miller index) and + # the derived-from-HKL fields below are recomputed, never gathered. + _REINDEX_FILL = { + "F": 0.0, + "I": 0.0, + "phase": 0.0, + "fom": 0.0, + "E": 0.0, + "E_squared": 0.0, + "F_squared_corrected": 0.0, + "radial_shell_indices": 0, + "F_sigma": 1.0, + "I_sigma": 1.0, + "rfree_flags": 1, # missing reflections default to the work set + "friedel_flags": False, + "validation_flags": False, + "outlier_flags": False, + } + + # Per-reflection fields that are pure functions of (hkl, cell, spacegroup): + # never gathered/aggregated, always recomputed or invalidated after an HKL + # change (``resolution`` recomputed; ``bin_indices`` / ``_centric_flags`` + # lazily rebuilt by ``get_bins`` / the ``centric`` property). + _REINDEX_DERIVED = ("resolution", "bin_indices", "_centric_flags") + + # Tensor dataclass fields that are NOT per-reflection (exempt from the + # length invariant): overall anisotropy parameters have shape (6,). + _NON_PER_REFLECTION_TENSORS = frozenset({"U_aniso"}) + + def _reindex_per_reflection( + self, + index_map: torch.Tensor, + new_hkl: torch.Tensor, + target: Optional["ReflectionData"] = None, + ) -> torch.Tensor: + """Gather every per-reflection field from ``self`` onto ``new_hkl``. + + Enumerates dataclass fields generically (same ``shape[0] == n`` guard as + :meth:`__select__` / :meth:`_canonicalize_in_place`) so any current or + future per-reflection field is carried automatically. + + Parameters + ---------- + index_map : torch.Tensor + LongTensor of length ``len(new_hkl)`` mapping each output row to a + source row index into ``self``, or ``-1`` for reflections absent + from the source (filled per :attr:`_REINDEX_FILL`). + new_hkl : torch.Tensor + Miller indices for the reindexed dataset, shape ``(M, 3)``. + target : ReflectionData, optional + Object to write the gathered fields onto. Defaults to ``self`` + (in-place). Pass a fresh instance for out-of-place reindexing + (``remap``); the caller is responsible for setting ``target.cell`` / + ``target.spacegroup`` before calling so resolution can be recomputed. + + Returns + ------- + torch.Tensor + Boolean presence mask (``index_map >= 0``), for the caller's use in + building ``hkl_present`` / ``missing`` masks. + """ + from dataclasses import fields as dc_fields + + if target is None: + target = self + + n_src = len(self.hkl) if self.hkl is not None else 0 + new_hkl = new_hkl.to(dtype=dtypes.int, device=self.device) + n_out = len(new_hkl) + index_map = index_map.to(device=self.device, dtype=torch.long) + present = index_map >= 0 + src_idx = index_map[present] + + derived = set(self._REINDEX_DERIVED) + for f in dc_fields(self): + name = f.name + if name == "hkl" or name in derived: + continue + # Declared non-per-reflection fields are exempt by *name*, not by + # shape. The shape test below is a heuristic and collides whenever + # n_src equals the field's own length -- ``U_aniso`` is (6,), so a + # 6-reflection dataset would have it gathered as if it were + # per-reflection. ``_assert_per_reflection_consistent`` and + # ``reduce_to_spacegroup`` already exempt by name; this keeps all + # three routines consistent. + if name in self._NON_PER_REFLECTION_TENSORS: + continue + val = getattr(self, name) + if not isinstance(val, torch.Tensor): + continue + if not (val.shape and val.shape[0] == n_src): + # Non-per-reflection tensor: leave target's own value untouched + # (a fresh default when target is a new instance). + continue + if name == "hkl_anomalous": + # Present rows keep their signed (anomalous) index; missing rows + # fall back to the canonical reference HKL (never a 0,0,0 row). + out = new_hkl.clone() + out[present] = val[src_idx] + else: + fill = self._REINDEX_FILL.get(name, 0) + out = torch.full( + (n_out,) + tuple(val.shape[1:]), + fill, + dtype=val.dtype, + device=self.device, + ) + out[present] = val[src_idx] + setattr(target, name, out) + + # Install the new HKL and recompute / invalidate derived-from-HKL fields. + target.hkl = new_hkl + target.bin_indices = None + target._centric_flags = None + if target.cell is not None: + target._calculate_resolution() + else: + target.resolution = None + return present + + def _assert_per_reflection_consistent(self) -> None: + """Invariant: every per-reflection tensor matches ``len(self.hkl)``. + + A cheap post-condition for the reindex routines. Turns any future + hardcoded-list regression (a field silently left at the old length) into + a loud failure instead of a downstream shape mismatch. + """ + from dataclasses import fields as dc_fields + + n = len(self.hkl) if self.hkl is not None else 0 + bad = [] + for f in dc_fields(self): + if f.name in self._NON_PER_REFLECTION_TENSORS: + continue + val = getattr(self, f.name) + if isinstance(val, torch.Tensor) and val.ndim >= 1 and val.shape[0] != n: + bad.append((f.name, tuple(val.shape))) + if bad: + raise RuntimeError( + f"ReflectionData per-reflection length mismatch (n_hkl={n}): {bad}" + ) + def _canonicalize_in_place(self) -> None: """Remap HKL to canonical CCP4 ASU form and reorder all data in-place.""" from dataclasses import fields as dc_fields @@ -455,7 +612,12 @@ def load(self, reader, french_wilson: bool = True): ) if spacegroup is not None: - self.spacegroup = SpaceGroup(spacegroup) + # Build on the dataset's device, not the process default: a + # ``ReflectionData(device='cpu')`` on an MPS host would otherwise + # end up with mps-resident symmetry matrices while every other + # tensor it owns is on CPU, and nothing downstream would notice + # until an op mixed the two. + self.spacegroup = SpaceGroup(spacegroup, device=self.device) self._calculate_resolution() # Prefer intensities (via French-Wilson) only when French-Wilson is @@ -640,9 +802,14 @@ def from_tensors( ReflectionData Fully initialized reflection data with all cleanup applied. """ - if device is None: - device = get_default_device() - data = cls(device=device, verbose=verbose) + # Follow the incoming tensors when no device is given, rather than + # the global default -- otherwise building from accelerator-resident + # arrays on a CPU-default host silently round-trips every one of them + # through host memory. ``hkl`` is a bare tensor, so read its device + # directly; ``resolve_device`` is for objects it can move in place. + if device is None and isinstance(hkl, torch.Tensor): + device = hkl.device + data = cls(device=normalize_device(device), verbose=verbose) def _prep(t: torch.Tensor) -> torch.Tensor: # Detach (constant data) unless the caller wants the graph preserved. @@ -657,7 +824,9 @@ def _prep(t: torch.Tensor) -> torch.Tensor: else Cell(cell, device=data.device) ) data.spacegroup = ( - spacegroup if isinstance(spacegroup, SpaceGroup) else SpaceGroup(spacegroup) + spacegroup + if isinstance(spacegroup, SpaceGroup) + else SpaceGroup(spacegroup, device=data.device) ) if rfree_flags is not None: @@ -691,7 +860,7 @@ def _prep(t: torch.Tensor) -> torch.Tensor: def load_mtz( self, - path: str, + path: Union[str, Path], column_names: Optional[dict] = None, french_wilson: bool = True, anomalous: Optional[bool] = None, @@ -723,14 +892,18 @@ def load_mtz( ReflectionData Self, for method chaining. """ + # ``str(path)``: gemmi's readers take a str, and the public + # ``torchref.io.read_mtz`` already advertises ``Union[str, Path]``, so + # accepting a Path here too keeps the two entry points consistent + # instead of failing deep inside gemmi with an argument-type error. reader = mtz.MTZReader( verbose=self.verbose, column_names=column_names, anomalous=anomalous - ).read(path) + ).read(str(path)) return self.load(reader, french_wilson=french_wilson) def load_cif( self, - path: str, + path: Union[str, Path], data_block: Optional[str] = None, anomalous: Optional[bool] = None, ) -> "ReflectionData": @@ -756,7 +929,7 @@ def load_cif( Self, for method chaining. """ self.reader = cif.ReflectionCIFReader( - path, verbose=self.verbose, data_block=data_block, anomalous=anomalous + str(path), verbose=self.verbose, data_block=data_block, anomalous=anomalous ) return self.load(self.reader) @@ -2069,7 +2242,12 @@ def __select__(self, indices: torch.Tensor, op=None) -> "ReflectionData": if val is None: continue if isinstance(val, torch.Tensor): - if val.shape and val.shape[0] == n_refl: + # Exempt by name first: the shape test is a heuristic that + # collides when n_refl equals the field's own length (see + # ``_reindex_per_reflection``). + if f.name in self._NON_PER_REFLECTION_TENSORS: + setattr(selected, f.name, val.clone()) + elif val.shape and val.shape[0] == n_refl: setattr(selected, f.name, val[indices]) else: # Non-matching tensor (e.g. U_aniso shape (6,)): copy as-is @@ -2223,63 +2401,15 @@ def validate_hkl(self, hkl_ref: torch.Tensor) -> "ReflectionData": [data_hkl_to_idx.get(tuple(hkl), -1) for hkl in hkl_ref_np], dtype=np.int64 ) - # Create presence mask: True where data exists - presence_mask = torch.from_numpy(ref_to_data_idx >= 0).to(device=self.device) + # Index map into this dataset for each reference row (-1 where missing). valid_indices = torch.from_numpy(ref_to_data_idx).to(device=self.device) - # Helper to expand a tensor to reference size - def expand_tensor(tensor, fill_value=0.0): - if tensor is None: - return None - expanded = torch.full( - (n_ref,) + tensor.shape[1:], - fill_value, - dtype=tensor.dtype, - device=self.device, - ) - # Copy existing data to correct positions - mask = valid_indices >= 0 - expanded[mask] = tensor[valid_indices[mask]] - return expanded - - # Expand all data arrays - old_F = self.F - old_F_sigma = self.F_sigma - old_I = self.I - old_I_sigma = getattr(self, "I_sigma", None) - old_rfree = self.rfree_flags - old_phase = getattr(self, "phase", None) - old_fom = getattr(self, "fom", None) - - # Replace HKL with reference - self.hkl = hkl_ref - - # Expand data tensors - self.F = expand_tensor(old_F, fill_value=0.0) - self.F_sigma = expand_tensor(old_F_sigma, fill_value=1.0) - - if old_I is not None: - self.I = expand_tensor(old_I, fill_value=0.0) - if old_I_sigma is not None: - self.I_sigma = expand_tensor(old_I_sigma, fill_value=1.0) - - # For rfree, default missing to work set (1) - if old_rfree is not None: - rfree_expanded = torch.ones( - n_ref, dtype=old_rfree.dtype, device=self.device - ) - mask = valid_indices >= 0 - rfree_expanded[mask] = old_rfree[valid_indices[mask]] - self.rfree_flags = rfree_expanded - - # Recalculate resolution for new HKL set - self._calculate_resolution() - - # Expand phase and fom if present - if old_phase is not None: - self.phase = expand_tensor(old_phase, fill_value=0.0) - if old_fom is not None: - self.fom = expand_tensor(old_fom, fill_value=0.0) + # Reindex EVERY per-reflection field onto the reference grid via the + # shared primitive (carries hkl_anomalous / friedel_flags / centric / + # validation / outlier that this routine used to silently drop, which + # broke difference refinement on mismatched reflection files). Masks are + # handled separately below because they are not dataclass fields. + presence_mask = self._reindex_per_reflection(valid_indices, hkl_ref) # Transfer existing masks to new indexing old_masks = dict(self.masks.items()) @@ -2308,6 +2438,7 @@ def expand_tensor(tensor, fill_value=0.0): print(f" Present in data: {n_present} ({100*n_present/n_ref:.1f}%)") print(f" Missing (masked): {n_missing} ({100*n_missing/n_ref:.1f}%)") + self._assert_per_reflection_consistent() return self def find_outliers( @@ -3046,8 +3177,9 @@ def remap( """ from torchref.symmetry.spacegroup import SpaceGroup - # Helper function for remapping tensors with missing handling - def _remap_tensor(tensor, fill_value): + # Mask remapper. Masks are not dataclass fields, so the shared + # per-reflection reindexer below does not touch them. + def _remap_mask(tensor, fill_value): if tensor is None: return None valid_mask = index_mapping >= 0 @@ -3060,56 +3192,36 @@ def _remap_tensor(tensor, fill_value): result[valid_mask] = tensor[index_mapping[valid_mask]] return result - # Create new ReflectionData + # Create new ReflectionData; set cell/spacegroup first so the shared + # reindexer can recompute resolution on the new grid. remapped = ReflectionData(verbose=self.verbose, device=self.device) + remapped.cell = self.cell.clone() if self.cell is not None else None + if spacegroup is not None: + remapped.spacegroup = SpaceGroup(spacegroup, device=remapped.device) + else: + remapped.spacegroup = self.spacegroup - # Set new HKL - remapped.hkl = new_hkl.to(device=self.device) - - # Remap amplitude and intensity fields - remapped.F = _remap_tensor(self.F, fill_value=0.0) - remapped.F_sigma = _remap_tensor(self.F_sigma, fill_value=1.0) - remapped.I = _remap_tensor(self.I, fill_value=0.0) - remapped.I_sigma = _remap_tensor(self.I_sigma, fill_value=1.0) - remapped.fom = _remap_tensor(self.fom, fill_value=0.0) - - # Carry forward prior mask if available - prior_mask = self.masks() - if prior_mask is not None: - remapped.masks["prior_flagged"] = _remap_tensor( - prior_mask.to(dtype=dtypes.int), fill_value=0 - ).to(torch.bool) - - # Handle rfree_flags (True = include in Rfree for missing) - if self.rfree_flags is not None: - remapped.rfree_flags = _remap_tensor( - self.rfree_flags.to(dtype=dtypes.int), fill_value=1 - ).to(torch.bool) + # Reindex ALL per-reflection dataclass fields onto new_hkl. This carries + # the anomalous / validation / centric / outlier fields the old hardcoded + # list silently dropped, and sets hkl / resolution while invalidating + # bin_indices and _centric_flags. Missing rows (index -1) get the + # per-field fills in _REINDEX_FILL. + self._reindex_per_reflection(index_mapping, new_hkl, target=remapped) - # Handle phases with optional shifts - if self.phase is not None: - remapped.phase = _remap_tensor(self.phase, fill_value=0.0) + # Apply optional phase shifts (e.g. from symmetry translations). + if remapped.phase is not None: if phase_shifts is not None: remapped.phase = remapped.phase + phase_shifts.to(device=self.device) elif phase_shifts is not None: - # Store phase shifts even if no original phases + # No original phases: store the shifts for later phase reconstruction. remapped._expansion_phase_shifts = phase_shifts.to(device=self.device) - # Clone cell - remapped.cell = self.cell.clone() if self.cell is not None else None - - # Set spacegroup - if spacegroup is not None: - remapped.spacegroup = SpaceGroup(spacegroup) - else: - remapped.spacegroup = self.spacegroup - - # Recalculate resolution - if remapped.cell is not None and remapped.hkl is not None: - remapped._calculate_resolution() - - # Invalidate dependent fields - remapped.bin_indices = None + # Carry forward prior combined mask if available. + prior_mask = self.masks() + if prior_mask is not None: + remapped.masks["prior_flagged"] = _remap_mask( + prior_mask.to(dtype=dtypes.int), fill_value=0 + ).to(torch.bool) # Copy metadata sources remapped.amplitude_source = self.amplitude_source @@ -3126,6 +3238,7 @@ def _remap_tensor(tensor, fill_value): if missing_mask.any(): remapped.masks["missing"] = ~missing_mask.to(device=self.device) + remapped._assert_per_reflection_consistent() return remapped def fill(self, d_min: Optional[float] = None) -> "ReflectionData": @@ -3440,18 +3553,69 @@ def _aggregate_sigma(tensor, agg_func="mean"): # 0 = free, non-zero = work. Take min to get free if any is free. reduced.rfree_flags = rfree_gathered.min(dim=1).values != 0 + # Boolean per-reflection flags: an equivalent's flag propagates to the + # merged reflection if ANY contributor has it set (validation/outlier + # are conservative "exclude if any"). + def _aggregate_any(tensor): + if tensor is None: + return None + gathered = tensor[reduction_indices.clamp(min=0)].to(torch.bool) + gathered = gathered & valid_mask + return gathered.any(dim=1) + + if self.validation_flags is not None: + reduced.validation_flags = _aggregate_any(self.validation_flags) + if self.outlier_flags is not None: + reduced.outlier_flags = _aggregate_any(self.outlier_flags) + + # Completeness pass: carry any remaining per-reflection dataclass tensor + # field not handled above so the merge never silently drops data (the + # historical bug). Derived-from-HKL fields are recomputed/invalidated + # below, not aggregated; 'first' is a safe representative for the rest. + from dataclasses import fields as dc_fields + + _already_set = { + "hkl", + "F", + "F_sigma", + "I", + "I_sigma", + "phase", + "fom", + "rfree_flags", + "validation_flags", + "outlier_flags", + } + _recomputed = set(self._REINDEX_DERIVED) | {"hkl_anomalous", "friedel_flags"} + n_src = len(self.hkl) + for f in dc_fields(self): + name = f.name + if name in _already_set or name in _recomputed: + continue + if name in self._NON_PER_REFLECTION_TENSORS: + continue + val = getattr(self, name) + if not isinstance(val, torch.Tensor): + continue + if not (val.shape and val.shape[0] == n_src): + continue + setattr(reduced, name, _aggregate_tensor(val, "first")) + # Clone cell reduced.cell = self.cell.clone() if self.cell is not None else None - # Set spacegroup - reduced.spacegroup = SpaceGroup(spacegroup) + # Set spacegroup (on the reduced dataset's device, not the global default) + reduced.spacegroup = SpaceGroup(spacegroup, device=reduced.device) # Recalculate resolution if reduced.cell is not None and reduced.hkl is not None: reduced._calculate_resolution() - # Invalidate dependent fields + # Invalidate derived-from-HKL fields (recomputed lazily for the new ASU). + # hkl_anomalous / friedel_flags are left unset: the merged ASU is + # Friedel-merged, so hkl_for_sf() correctly falls back to hkl. reduced.bin_indices = None + reduced._centric_flags = None # Copy metadata sources reduced.amplitude_source = self.amplitude_source @@ -3465,6 +3629,7 @@ def _aggregate_sigma(tensor, agg_func="mean"): f"reduce_to_spacegroup({spacegroup}, aggregation={aggregation})" ) + reduced._assert_per_reflection_consistent() return reduced def canonicalize(self, include_friedel: bool = True) -> "ReflectionData": diff --git a/torchref/io/ihm.py b/torchref/io/ihm.py index 5b1e858c..a054dc68 100644 --- a/torchref/io/ihm.py +++ b/torchref/io/ihm.py @@ -474,10 +474,9 @@ def build_model_collection( "Call read_atom_data() first." ) - if device is None: - from torchref.config import get_default_device + from torchref.config import normalize_device - device = get_default_device() + device = normalize_device(device) # Build one ModelFT per state base_models = [] diff --git a/torchref/maps/difference_map.py b/torchref/maps/difference_map.py index 44a19e1a..6051c5ed 100644 --- a/torchref/maps/difference_map.py +++ b/torchref/maps/difference_map.py @@ -12,7 +12,6 @@ import torch from torchref.base.reciprocal.grid_operations import place_on_grid -from torchref.config import get_default_device from torchref.io.datasets.collection import DatasetCollection from torchref.maps.map import Map from torchref.symmetry.reciprocal_symmetry import expand_hkl diff --git a/torchref/model/mixed_model.py b/torchref/model/mixed_model.py index 787cebe2..e24e443e 100644 --- a/torchref/model/mixed_model.py +++ b/torchref/model/mixed_model.py @@ -12,6 +12,7 @@ from torch import nn from torchref.utils.device_mixin import DeviceMovementMixin +from torchref.utils.device_resolution import resolve_device if TYPE_CHECKING: from torchref.model.model_ft import ModelFT @@ -104,12 +105,12 @@ def __init__( self.verbose = verbose - # Infer device from first model if not specified - if device is None: - device = models[0].device + # First-wins reconciliation across the supplied models, and the move + # itself: ``resolve_device`` warns when they disagree, which the manual + # loop this replaces did silently. + device = resolve_device(*models, device=device) # Store models as ModuleList for proper PyTorch handling - models = [model.to(device=device) for model in models] self.models = nn.ModuleList(models) # Validate model compatibility diff --git a/torchref/model/model.py b/torchref/model/model.py index 94d4c00d..6056a1df 100644 --- a/torchref/model/model.py +++ b/torchref/model/model.py @@ -19,7 +19,7 @@ import torch.nn as nn from torchref.base import math_torch -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_float_dtype, normalize_device from torchref.io import cif, pdb from torchref.model.parameter_wrappers import ( CholeskyMixedTensor, @@ -151,8 +151,7 @@ def __init__( # ``dtypes.float`` / ``device.current`` change is honored. if dtype_float is None: dtype_float = get_float_dtype() - if device is None: - device = get_default_device() + device = normalize_device(device) # Configuration self.dtype_float = dtype_float self.verbose = verbose @@ -287,7 +286,10 @@ def spacegroup(self, value): a space group name string, or a space group number. """ if value is not None: - self._spacegroup = SpaceGroup(value) + # ``device=self.device``: SpaceGroup falls back to the global + # default otherwise, so setting a spacegroup on a CPU-pinned Model + # would silently plant accelerator-resident matrices on it. + self._spacegroup = SpaceGroup(value, device=self.device) else: self._spacegroup = None @@ -1091,21 +1093,28 @@ def get_vdw_radii(self): ), f"vdW radii length mismatch with number of atoms {len(self.vdw_radii)} != {len(self.pdb)}" return self.vdw_radii - def to(self, *args, **kwargs): - """Move Model and rebuild device-specific SF indices. + def _after_device_apply( + self, old_device, new_device, old_dtype, new_dtype, *, + device_changed, dtype_changed, + ): + """Rebuild the precomputed SF indices after a real device/dtype change. - Delegates to :class:`~torchref.utils.device_mixin.DeviceMovementMixin`, which - walks ``self.__dict__`` (picking up ``self.cell``, ``self.altloc_pairs``, + :class:`~torchref.utils.device_mixin.DeviceMixin` walks + ``self.__dict__`` (picking up ``self.cell``, ``self.altloc_pairs``, ``self._restraints`` and all registered parameters / buffers), refreshes - the ``self.device`` tracker, and invalidates caches. Afterwards this - override rebuilds the precomputed SF indices on the new device. - """ - result = super().to(*args, **kwargs) - if hasattr(result, "aniso_flag") and result.aniso_flag is not None: - result._rebuild_sf_indices() - if result.verbose > 0: - print(f"Model moved to device: {result.device}") - return result + the ``self.device`` tracker and invalidates caches; this hook then + regenerates the iso/aniso index tensors on the new device. + + This lives on the movement hook rather than a ``to()`` override because + a parent module reaches ``Model`` through ``_apply``, which never calls + ``to()``. It is deliberately not ``reset_cache()``: that hook fires + after every optimizer step (see ``LossState.reset_caches``), and index + reconstruction does not belong on that path. + """ + if getattr(self, "aniso_flag", None) is not None: + self._rebuild_sf_indices() + if self.verbose > 0: + print(f"Model moved to device: {self.device}") def copy(self): """ @@ -2934,8 +2943,7 @@ def create_from_state_dict( """ # Resolve dtype/device at call time so the fallbacks below use the # current config, not the import-time default. - if device is None: - device = get_default_device() + device = normalize_device(device) if dtype_float is None: dtype_float = get_float_dtype() # Extract metadata (non-tensor data that we handle specially) diff --git a/torchref/model/model_collection.py b/torchref/model/model_collection.py index 34084c2f..c798872a 100644 --- a/torchref/model/model_collection.py +++ b/torchref/model/model_collection.py @@ -11,8 +11,8 @@ import torch from torch import nn -from torchref.config import get_default_device from torchref.utils.device_mixin import DeviceMovementMixin +from torchref.utils.device_resolution import resolve_device if TYPE_CHECKING: from torchref.model.model_ft import ModelFT @@ -69,8 +69,11 @@ def __init__( # Normalize to handle floating point drift initial_fractions = [f / total for f in initial_fractions] - if device is None: - device = base_models[0].device if hasattr(base_models[0], "device") else get_default_device() + # Reconcile across *all* base models, not just the first: reading + # ``base_models[0].device`` left a mixed-device list unreconciled, so + # ``fractions_tensor`` below could land on a device the later models + # were not on. + device = resolve_device(*base_models, device=device) # Match base models' float dtype (consistent under a float64 config). fractions_tensor = torch.tensor( diff --git a/torchref/model/model_ft.py b/torchref/model/model_ft.py index 214fc0c6..a4195edb 100644 --- a/torchref/model/model_ft.py +++ b/torchref/model/model_ft.py @@ -6,7 +6,7 @@ import torch from torchref.base.fourier import fft, ifft -from torchref.config import dtypes, get_default_device, get_float_dtype +from torchref.config import dtypes, get_float_dtype, normalize_device from torchref.model.model import Model from torchref.model.sf_fft import SfFFT from torchref.symmetry import SpaceGroup @@ -1151,8 +1151,7 @@ def create_from_state_dict( # Resolve dtype/device at call time so the fallback below uses the # current config rather than an import-time default. - if device is None: - device = get_default_device() + device = normalize_device(device) if dtype_float is None: dtype_float = get_float_dtype() diff --git a/torchref/model/parameter_wrappers.py b/torchref/model/parameter_wrappers.py index 50289f8d..550d68c6 100644 --- a/torchref/model/parameter_wrappers.py +++ b/torchref/model/parameter_wrappers.py @@ -8,6 +8,7 @@ import torch from torch import nn +from torchref.config import get_float_dtype, normalize_device from torchref.utils.caching import CachedForwardMixin from torchref.utils.device_mixin import DeviceMixin @@ -133,12 +134,19 @@ def __init__( # Empty initialization if initial_values is None: + # Honour the requested device/dtype even with no values yet: the + # empty shell is the documented ``load_state_dict`` entry point, and + # a caller that asked for a device should not get a CPU parameter + # (nor a ``.device`` property that raises because every buffer is + # ``None``). + device = normalize_device(device) + dtype = dtype if dtype is not None else get_float_dtype() self.register_buffer("refinable_mask", None) self.register_buffer("fixed_mask", None) self.register_buffer("fixed_values", None) - self.register_buffer("_shape", None) self.refinable_params = nn.Parameter( - torch.empty(0), requires_grad=requires_grad + torch.empty(0, device=device, dtype=dtype), + requires_grad=requires_grad, ) self._has_refinable = False self._refinable_indices = None @@ -200,8 +208,12 @@ def __init__( refinable_values, requires_grad=requires_grad ) - # Store shape for reconstruction - self.register_buffer("_shape", torch.tensor(initial_values.shape)) + # Shape is host-side metadata, deliberately *not* a registered buffer: + # a buffer would be dragged onto the accelerator by ``.to()`` and then + # read back with ``.tolist()``, i.e. a device sync on every ``.shape`` + # access, purely to recover numbers we already have. It is also fully + # derivable from ``fixed_values`` (a clone of ``initial_values``), so + # storing it separately only creates something that can disagree. # Pre-compute index cache to avoid boolean indexing at runtime self._build_index_cache() @@ -476,17 +488,30 @@ def set(self, values: torch.Tensor, mask: torch.Tensor) -> None: @property def shape(self): """Return the shape of the full tensor.""" - return tuple(self._shape.tolist()) + if self.fixed_values is None: + return () + return tuple(self.fixed_values.shape) + + # ``fixed_values`` is ``None`` on the empty-shell path, so both properties + # fall back to the (empty) parameter, which always exists and carries the + # device/dtype the constructor was given. Without the fallback the + # ``AttributeError`` raised by ``None.device`` is intercepted by + # ``nn.Module.__getattr__`` and re-raised as the thoroughly misleading + # "'MixedTensor' object has no attribute 'device'". @property def dtype(self): """Return the dtype of the tensor.""" - return self.fixed_values.dtype + if self.fixed_values is not None: + return self.fixed_values.dtype + return self.refinable_params.dtype @property def device(self): """Return the device of the tensor.""" - return self.fixed_values.device + if self.fixed_values is not None: + return self.fixed_values.device + return self.refinable_params.device def get_refinable_count(self) -> int: """Return the number of refinable parameters.""" @@ -517,6 +542,16 @@ def update_fixed_values(self, new_values: torch.Tensor): ) self.fixed_values = new_values.to(dtype=self.dtype, device=self.device).detach() + def _normalize_refinable_mask(self, new_mask: torch.Tensor) -> torch.Tensor: + """Coerce an incoming mask onto this wrapper's device as ``bool``. + + Must be applied *before* deriving ``fixed_mask = ~new_mask``: negating + the caller's un-migrated tensor leaves ``refinable_mask`` and + ``fixed_mask`` on different devices, which then fails far away from + here, in whichever op first uses both. + """ + return new_mask.to(device=self.device, dtype=torch.bool) + def update_refinable_mask( self, new_mask: torch.Tensor, reset_refinable: bool = False ): @@ -542,7 +577,8 @@ def update_refinable_mask( current_full = self.forward().detach() - self.refinable_mask = new_mask.to(device=self.device) + new_mask = self._normalize_refinable_mask(new_mask) + self.refinable_mask = new_mask self.fixed_mask = ~new_mask if reset_refinable: @@ -609,11 +645,27 @@ def clip(self, min_value=None, max_value=None) -> "MixedTensor": ) return new_mixed - def to(self, *args, **kwargs): - """Move via :class:`DeviceMixin` and rebuild the index cache.""" - result = super().to(*args, **kwargs) - result._build_index_cache() - return result + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + """Drop the legacy ``_shape`` buffer key before delegating. + + ``_shape`` used to be a registered buffer; it is now derived from + ``fixed_values``. Checkpoints written before that change still carry the + key, and a ``strict=True`` load would reject it as unexpected. + """ + state_dict.pop(prefix + "_shape", None) + return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + + def _after_device_apply(self, *args, device_changed, dtype_changed, **kwargs): + """Rebuild the index cache after a real device/dtype change. + + ``_refinable_indices`` / ``_refinable_idx_1d`` are plain attributes + holding tensors, so they must be regenerated on the new device. Using + the movement hook rather than a ``to()`` override matters twice over: + ``to()`` is bypassed when a parent module reaches this wrapper through + ``_apply``, and ``reset_cache()`` would put this rebuild on the + per-optimizer-step path. + """ + self._build_index_cache() def refine( self, selection: Union[slice, torch.Tensor, tuple], reset_values: bool = False @@ -676,6 +728,7 @@ def refine( new_mask |= temp_mask # Update the mask + new_mask = self._normalize_refinable_mask(new_mask) self.refinable_mask = new_mask self.fixed_mask = ~new_mask @@ -755,6 +808,7 @@ def fix( new_mask &= ~temp_mask # Update the mask + new_mask = self._normalize_refinable_mask(new_mask) self.refinable_mask = new_mask self.fixed_mask = ~new_mask @@ -1180,7 +1234,8 @@ def update_refinable_mask( # Convert to log space current_log = torch.log(current_normal.clamp(min=self.epsilon)) - self.refinable_mask = new_mask.to(device=self.device) + new_mask = self._normalize_refinable_mask(new_mask) + self.refinable_mask = new_mask self.fixed_mask = ~new_mask # Update fixed_values with log-space values @@ -1407,7 +1462,8 @@ def update_refinable_mask( ) with torch.no_grad(): current_raw = self._u6_to_raw6(self.forward()) - self.refinable_mask = new_mask.to(device=self.device) + new_mask = self._normalize_refinable_mask(new_mask) + self.refinable_mask = new_mask self.fixed_mask = ~new_mask self.fixed_values = current_raw.clone() new_refinable = current_raw[self.refinable_mask].clone() @@ -1556,16 +1612,22 @@ def __init__( self._name = name or "occupancy" # Empty initialization + # + # ``OccupancyTensor`` builds its own shell rather than delegating to + # ``MixedTensor``, so the device/dtype handling has to be repeated here + # -- fixing only the base class would leave this path on CPU. if initial_values is None: + device = normalize_device(device) + dtype = dtype if dtype is not None else get_float_dtype() self._full_shape = 0 self._collapsed_shape = 0 self.register_buffer("refinable_mask", None) self.register_buffer("fixed_mask", None) self.register_buffer("fixed_values", None) - self.register_buffer("_shape", None) self.register_buffer("expansion_mask", None) self.refinable_params = nn.Parameter( - torch.empty(0), requires_grad=requires_grad + torch.empty(0, device=device, dtype=dtype), + requires_grad=requires_grad, ) self._has_refinable = False self._refinable_indices = None @@ -1644,8 +1706,11 @@ def __init__( refinable_values, requires_grad=requires_grad ) - # Store collapsed shape - self.register_buffer("_shape", torch.tensor([self._collapsed_shape])) + # Collapsed shape is host-side metadata, not a buffer -- same reasoning + # as ``MixedTensor`` (a buffer gets dragged onto the accelerator by + # ``.to()`` and then read back with a sync, purely to recover a number + # we already have). ``_collapsed_shape`` already holds it as an int; + # the ``shape`` property below derives from it. # Pre-compute index cache to avoid boolean indexing at runtime self._build_index_cache() diff --git a/torchref/model/rigid_xyz.py b/torchref/model/rigid_xyz.py index e99f56ed..1c40feb4 100644 --- a/torchref/model/rigid_xyz.py +++ b/torchref/model/rigid_xyz.py @@ -18,6 +18,7 @@ from torch import nn from torchref.base.alignment.rotation import rotation_matrix_euler_xyz +from torchref.config import get_float_dtype, normalize_device from torchref.utils.caching import CachedForwardMixin from torchref.utils.device_mixin import DeviceMixin @@ -56,14 +57,22 @@ def __init__( self._name = name if original_xyz is None: - # Empty init for state_dict loading. - self.register_buffer("original_xyz", torch.empty(0, 3)) - self.register_buffer("chain_indices", torch.empty(0, dtype=torch.long)) - self.register_buffer("chain_centers", torch.empty(0, 3)) - self.register_buffer("mobile_mask", torch.empty(0, dtype=torch.bool)) - self.register_buffer("atom_weights", torch.empty(0)) - self.euler_angles = nn.Parameter(torch.empty(0, 3)) - self.translations = nn.Parameter(torch.empty(0, 3)) + # Empty init for state_dict loading. Honour the requested + # device/dtype: otherwise every buffer lands on CPU regardless of + # what the caller asked for, and only a later ``.to()`` repairs it. + device = normalize_device(device) + dtype = dtype if dtype is not None else get_float_dtype() + self.register_buffer("original_xyz", torch.empty(0, 3, device=device, dtype=dtype)) + self.register_buffer( + "chain_indices", torch.empty(0, dtype=torch.long, device=device) + ) + self.register_buffer("chain_centers", torch.empty(0, 3, device=device, dtype=dtype)) + self.register_buffer( + "mobile_mask", torch.empty(0, dtype=torch.bool, device=device) + ) + self.register_buffer("atom_weights", torch.empty(0, device=device, dtype=dtype)) + self.euler_angles = nn.Parameter(torch.empty(0, 3, device=device, dtype=dtype)) + self.translations = nn.Parameter(torch.empty(0, 3, device=device, dtype=dtype)) self._n_chains = 0 self._chain_id_order: list = [] return diff --git a/torchref/model/sf_ds.py b/torchref/model/sf_ds.py index e122c5fd..eacbdaf0 100644 --- a/torchref/model/sf_ds.py +++ b/torchref/model/sf_ds.py @@ -16,7 +16,6 @@ import torch.nn as nn from torchref.base.direct_summation import ( - Engine, ds_aniso, ds_iso, ) @@ -101,7 +100,7 @@ def __init__( device: torch.device = None, verbose: int = 0, max_memory_gb: float = 2.0, - engine: Engine = Engine.AUTO, + force_portable: Optional[bool] = None, ): """ Initialize the SfDS module with cell and spacegroup. @@ -120,29 +119,35 @@ def __init__( Verbosity level for logging. Default is 0. max_memory_gb : float, optional Maximum memory for intermediate tensors in GB. Default is 2.0. - engine : Engine, optional - Structure-factor backend selector. ``Engine.AUTO`` (default) - derives the backend from device/dtype/availability (Triton on - CUDA+float32, else checkpointed eager). ``Engine.TRITON`` and - ``Engine.EAGER`` force a path (for tests/benchmarks). + force_portable : bool, optional + Pin the portable reference path instead of the fastest usable backend, applied + per call so two instances can differ within one process. ``None`` (default) + defers to the process-wide setting, so ``with use_portable():`` steers an + unconfigured instance. + + ``SfFFT`` deliberately has no equivalent -- it never had a backend selector, and + ``use_portable()`` covers the same need by scoping the call. """ super().__init__() if dtype_float is None: dtype_float = dtypes.float - if device is None: - device = get_default_device() self.dtype_float = dtype_float - self.device = device + # Derive from ``cell`` when no device is given, instead of jumping to + # the global default and leaving a caller-supplied cell behind on + # another device. An explicit ``device`` still wins and moves the cell. + self.device = resolve_device(cell, device=device) self.verbose = verbose self.max_memory_gb = max_memory_gb - self.engine = engine + self.force_portable = force_portable # Store cell and spacegroup self._cell = cell self._spacegroup = None if spacegroup is not None or cell is not None: - self._spacegroup = SpaceGroup(spacegroup, dtype=dtype_float, device=device) + self._spacegroup = SpaceGroup( + spacegroup, dtype=dtype_float, device=self.device + ) # Cache reciprocal basis matrix self._recB: Optional[torch.Tensor] = None @@ -201,7 +206,13 @@ def set_cell_and_spacegroup(self, cell: Cell, spacegroup: SpaceGroupLike = None) Unit cell object. spacegroup : SpaceGroupLike, optional Space group specification. + + Notes + ----- + Receiver wins: an incoming cell on another device is moved to match + this module rather than the other way round. """ + self.device = resolve_device(self, cell) self._cell = cell self._recB = None # Invalidate cache self.spacegroup = spacegroup @@ -522,7 +533,7 @@ def _compute_p1_sf( adp_iso, A_iso, B_iso, - engine=self.engine, + force_portable=self.force_portable, max_memory_gb=self.max_memory_gb, ) sf_total = sf_total + sf_iso.to(sf_total.dtype) @@ -537,7 +548,7 @@ def _compute_p1_sf( u_aniso, A_aniso, B_aniso, - engine=self.engine, + force_portable=self.force_portable, max_memory_gb=self.max_memory_gb, ) sf_total = sf_total + sf_aniso.to(sf_total.dtype) @@ -576,7 +587,7 @@ def copy(self) -> "SfDS": device=self.device, verbose=self.verbose, max_memory_gb=self.max_memory_gb, - engine=self.engine, + force_portable=self.force_portable, ) return new_ds diff --git a/torchref/model/sf_fft.py b/torchref/model/sf_fft.py index 182f6de7..d2deb8fb 100644 --- a/torchref/model/sf_fft.py +++ b/torchref/model/sf_fft.py @@ -25,6 +25,7 @@ from torchref.symmetry.spacegroup import SpaceGroupLike from torchref.utils.caching import ParameterFingerprint from torchref.utils.device_mixin import DeviceMovementMixin +from torchref.utils.device_resolution import resolve_device class SfFFT(DeviceMovementMixin, nn.Module): @@ -127,11 +128,10 @@ def __init__( dtype_float = dtypes.float self.dtype_float = dtype_float - self.device = ( - device - if device is not None - else cell.device if cell is not None else get_default_device() - ) + # One device for the module and everything it builds. ``resolve_device`` + # also moves ``cell`` when an explicit ``device`` disagrees with it, so + # the cell and the SpaceGroup below cannot end up split. + self.device = resolve_device(cell, device=device) self.verbose = verbose self.use_late_symmetry = use_late_symmetry @@ -141,7 +141,12 @@ def __init__( self._spacegroup = None if spacegroup is not None or cell is not None: - self._spacegroup = SpaceGroup(spacegroup, dtype=dtype_float, device=device) + # ``self.device``, not the raw ``device`` argument: the latter is + # ``None`` on the derive-from-cell path, which would silently put + # the symmetry matrices on the global default instead. + self._spacegroup = SpaceGroup( + spacegroup, dtype=dtype_float, device=self.device + ) # Buffers (registered during setup_grid) self.register_buffer("gridsize", None) @@ -227,7 +232,14 @@ def set_cell_and_spacegroup(self, cell: Cell, spacegroup: SpaceGroupLike = None) Unit cell object. spacegroup : SpaceGroupLike, optional Space group specification. + + Notes + ----- + Receiver wins: this module may already own grid buffers, so an incoming + cell on another device is moved to match rather than dragging the + module after it. """ + self.device = resolve_device(self, cell) self._cell = cell self.spacegroup = spacegroup @@ -305,8 +317,10 @@ def compute_real_space_grid( torch.Tensor Real-space grid with shape (nx, ny, nz, 3). """ - if device is None: - device = get_default_device() + # Forward ``device`` as-is, including ``None``. ``get_real_grid`` + # already infers from ``fractional_matrix`` when no device is given; + # resolving the global default here meant it never saw ``None`` and + # its inference was dead code. return get_real_grid( fractional_matrix=fractional_matrix, gridsize=gridsize, device=device ) diff --git a/torchref/refinement/base_refinement.py b/torchref/refinement/base_refinement.py index 79c21ab5..bfff2886 100644 --- a/torchref/refinement/base_refinement.py +++ b/torchref/refinement/base_refinement.py @@ -7,7 +7,7 @@ import torch from torch.nn import Module as nnModule -from torchref.config import get_default_device +from torchref.config import normalize_device from torchref.io import ReflectionData from torchref.model.model_ft import ModelFT from torchref.refinement.logger import Logger @@ -1207,8 +1207,7 @@ def create_from_state_dict( print(f"Restored at R-work={rwork:.4f}, R-free={rfree:.4f}") """ - if device is None: - device = get_default_device() + device = normalize_device(device) # Helper to extract submodule state from flattened state_dict def extract_submodule_state(state_dict: dict, prefix: str) -> dict: diff --git a/torchref/refinement/loss_state.py b/torchref/refinement/loss_state.py index dad74447..4fc5308c 100644 --- a/torchref/refinement/loss_state.py +++ b/torchref/refinement/loss_state.py @@ -31,7 +31,7 @@ import torch from torch import nn -from torchref.config import get_default_device +from torchref.config import canonical_device, get_default_device from torchref.utils.autograd_introspection import collect_loss_leaves, _iter_roots from torchref.utils.device_mixin import DeviceMovementMixin from torchref.utils.loss_validation import validate_loss @@ -276,8 +276,33 @@ def register_target( if probe: self._probe_and_merge_leaves(target) self._collect_resettable_modules(target) + self._warn_on_device_mismatch(key, target) return self + def _warn_on_device_mismatch(self, key: str, target: Callable) -> None: + """Warn if ``target`` does not agree with this state's device. + + Deliberately a warning and not a move. ``target.to(self.device)`` would + drag the model, data and scaler the target *borrows* along with it, so + registering a loss term could silently relocate an entire + ``ReflectionData``. Targets resolve their own device at construction + (see ``Target._adopt_device``); a disagreement here means something + upstream built them inconsistently, which is worth surfacing rather + than papering over. + """ + target_device = getattr(target, "device", None) + if target_device is None or not isinstance(self.device, torch.device): + return + if canonical_device(target_device) == canonical_device(self.device): + return + warnings.warn( + f"LossState: target {key!r} is on {target_device} but the state is " + f"on {self.device}. Not moving it -- a target borrows its model and " + "data, so relocating it here would move those too. Construct the " + "target on the right device instead.", + stacklevel=3, + ) + def _collect_resettable_modules(self, target: Callable) -> None: """Walk ``target``'s submodules and collect any that expose a ``reset_cache`` method, deduplicating against modules already @@ -914,23 +939,12 @@ def summary(self) -> None: # ========================================================================= # Device Management # ========================================================================= - - def to(self, *args, **kwargs): - """Move via :class:`DeviceMixin`; honour an explicit device when no tensors exist yet.""" - result = super().to(*args, **kwargs) - # If no tensor was found to refresh ``self.device``, fall back to the - # explicit device argument so subsequent allocations land correctly. - if not isinstance(result.device, torch.device): - from torchref.utils.device_mixin import _parse_to_args - - device, _ = _parse_to_args(args, kwargs) - if device is not None: - result.device = ( - torch.device(device) - if not isinstance(device, torch.device) - else device - ) - return result + # + # ``LossState`` normally owns no tensors (targets are callables, weights are + # floats), so its ``device`` tracker used to need a bespoke ``to()`` + # override to survive a move. ``DeviceMixin`` now carries the parsed request + # down the traversal and updates tensor-free objects itself, so the override + # is gone; ``self.device`` follows ``.to()`` via the shared machinery. # ========================================================================= # Utility diff --git a/torchref/refinement/optimizers/__init__.py b/torchref/refinement/optimizers/__init__.py index 6e202c70..5cf1a3bc 100644 --- a/torchref/refinement/optimizers/__init__.py +++ b/torchref/refinement/optimizers/__init__.py @@ -6,12 +6,19 @@ - ExploratoryLBFGS: LBFGS with automatic landscape exploration via Lanczos - LangevinSA: BAOAB Langevin dynamics with simulated annealing - MomentumStochasticSA: Phenix-style SA (gradient descent + momentum + noise) +- SeededLBFGS: L-BFGS whose first step is a diagonal-Newton step """ from .adam_noise import AdamWithAdaptiveNoise +from .curvature import ( + hessian_diagonal, + hessian_diagonal_preconditioner, + preconditioner_from_diagonal, +) from .exploratory_lbfgs import ExploratoryLBFGS from .langevin_sa import LangevinSA from .momentum_sa import MomentumStochasticSA +from .seeded_lbfgs import SeededLBFGS __all__ = [ @@ -19,4 +26,8 @@ "ExploratoryLBFGS", "LangevinSA", "MomentumStochasticSA", + "SeededLBFGS", + "hessian_diagonal", + "hessian_diagonal_preconditioner", + "preconditioner_from_diagonal", ] diff --git a/torchref/refinement/optimizers/curvature.py b/torchref/refinement/optimizers/curvature.py new file mode 100644 index 00000000..ec6a2d3d --- /dev/null +++ b/torchref/refinement/optimizers/curvature.py @@ -0,0 +1,256 @@ +"""Hessian-diagonal estimation for seeding :class:`SeededLBFGS`. + +Provides a stochastic (Hutchinson) estimate of the diagonal of the loss Hessian +and the derived inverse-curvature preconditioner ``1/(|diag|+lambda)`` used to +seed the first L-BFGS step. + +The estimate uses *analytic* double-backward Hessian-vector products:: + + (g,) = torch.autograd.grad(loss, params, create_graph=True) # once + (Hv,) = torch.autograd.grad((g * v).sum(), params) # per probe + +which is only correct on the pure-torch electron-density path, so everything runs +under ``use_portable()``. Finite-difference HVPs are deliberately NOT +used: the production Triton/CUDA kernels emit float32 gradients, so central +differences ``(g(x+eps v) - g(x-eps v))`` catastrophically cancel. The analytic +HVP has no such cancellation and float32 is adequate for a preconditioner. + +The returned diagonal is flat, aligned to ``cat(p.reshape(-1) for p in params)`` +-- the same layout as ``torch.optim.LBFGS._gather_flat_grad`` -- so it can be fed +straight to :meth:`SeededLBFGS.set_init_hess_diag`. +""" + +from __future__ import annotations + +from typing import Callable, Optional, Sequence + +import torch + +from torchref.utils import use_portable + + +def _sample_probe( + numel: int, + probe: str, + generator: Optional[torch.Generator], + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Draw one Hutchinson probe vector of length ``numel``.""" + if probe == "rademacher": + r = torch.randint( + 0, 2, (numel,), generator=generator, device=device, dtype=torch.int64 + ) + return r.to(dtype).mul_(2.0).sub_(1.0) # {0,1} -> {-1,+1} + if probe == "gaussian": + return torch.randn(numel, generator=generator, device=device, dtype=dtype) + raise ValueError(f"unknown probe type {probe!r}") + + +def _flatten(tensors: Sequence[torch.Tensor]) -> torch.Tensor: + return torch.cat([t.reshape(-1) for t in tensors]) + + +def hessian_diagonal( + aggregate: Callable[[], torch.Tensor], + params: Sequence[torch.nn.Parameter], + *, + n_probes: int = 16, + generator: Optional[torch.Generator] = None, + probe: str = "rademacher", + basis_max_numel: int = 4096, +) -> torch.Tensor: + """Estimate the diagonal of the Hessian of ``aggregate()`` w.r.t. ``params``. + + Parameters + ---------- + aggregate : callable + Zero-argument callable returning the scalar loss tensor (e.g. + ``LossState.aggregate``). Called once; the graph is retained so that + ``n_probes`` Hessian-vector products can reuse the first-order backward. + params : sequence of nn.Parameter + The leaves to differentiate w.r.t., in the SAME order they are passed to + the optimizer (so the flat result matches ``_gather_flat_grad``). + n_probes : int + Number of Hutchinson probes (ignored when ``probe="basis"``). + generator : torch.Generator, optional + For deterministic probes. + probe : {"rademacher", "gaussian", "basis"} + ``"basis"`` computes the EXACT diagonal in ``numel`` HVPs (verification / + small problems only; capped by ``basis_max_numel``). + + Notes + ----- + The double-backward runs under :func:`torchref.utils.use_portable`, which + pins the portable reference kernel -- the only path that composes correctly + under ``create_graph=True`` for the Hessian-vector products. + + Returns + ------- + torch.Tensor + Flat detached estimate of ``diag(H)`` aligned to + ``cat(p.reshape(-1) for p in params)``. + """ + params = list(params) + if any(torch.is_complex(p) for p in params): + raise ValueError("hessian_diagonal does not support complex parameters") + # Differentiate only through leaves that require grad; frozen leaves (e.g. + # occupancy held fixed, or aniso U on an isotropic model) get zero curvature + # -> a bounded 1/lam seed. Harmless: their gradient is zero, so the optimizer + # never moves them regardless of the seed value. The flat layout still spans + # the FULL params list so it matches the optimizer's _gather_flat_grad order. + active = [p for p in params if p.requires_grad] + if not active: + raise ValueError("no params require grad; nothing to differentiate") + + def _assemble_full(active_parts): + """Map a per-active-param list back to a full-length flat tensor in + ``params`` order, filling detached zeros for frozen params.""" + it = iter(active_parts) + pieces = [] + for p in params: + if p.requires_grad: + pieces.append(next(it).reshape(-1)) + else: + pieces.append(torch.zeros(p.numel(), device=p.device, dtype=p.dtype)) + return torch.cat(pieces) + + with use_portable(): + loss = aggregate() + # allow_unused: an active leaf may still not enter the loss. + grads = torch.autograd.grad( + loss, active, create_graph=True, allow_unused=True + ) + grads = [ + g if g is not None else torch.zeros_like(p) + for g, p in zip(grads, active) + ] + g_flat = _assemble_full(grads) # graph-attached for active slices + numel = g_flat.numel() + device, dtype = g_flat.device, g_flat.dtype + + def _hvp(v: torch.Tensor, retain: bool) -> torch.Tensor: + parts = torch.autograd.grad( + (g_flat * v).sum(), + active, + retain_graph=retain, + create_graph=False, + allow_unused=True, + ) + parts = [ + pt if pt is not None else torch.zeros_like(p) + for pt, p in zip(parts, active) + ] + return _assemble_full(parts) + + diag = torch.zeros(numel, device=device, dtype=dtype) + + if probe == "basis": + if numel > basis_max_numel: + raise ValueError( + f"probe='basis' needs {numel} HVPs; exceeds basis_max_numel " + f"({basis_max_numel}). Use rademacher/gaussian for large problems." + ) + for i in range(numel): + v = torch.zeros(numel, device=device, dtype=dtype) + v[i] = 1.0 + diag[i] = _hvp(v, retain=(i < numel - 1))[i] + return _sanitize_diag(diag.detach()) + + for k in range(n_probes): + v = _sample_probe(numel, probe, generator, device, dtype) + diag.add_(v * _hvp(v, retain=(k < n_probes - 1))) + diag.div_(n_probes) + return _sanitize_diag(diag.detach()) + + +def _sanitize_diag(diag: torch.Tensor) -> torch.Tensor: + """Replace sporadic non-finite diagonal entries with the max finite |diag| + (→ stiffest, i.e. smallest, safest seeded step for those directions). If NO + entry is finite (the whole HVP blew up), leave it non-finite so the caller + can detect it and skip seeding this cycle. + """ + finite = torch.isfinite(diag) + if finite.all() or not finite.any(): + return diag + fill = diag[finite].abs().max() + return torch.where(finite, diag, fill) + + +def preconditioner_from_diagonal( + diag: torch.Tensor, lam: float = 1e-2, group_sizes=None +) -> torch.Tensor: + """Turn a Hessian diagonal into a positive inverse-curvature preconditioner. + + ``precond = 1 / max(|diag|, lam * max|diag|)`` -- the absolute value makes + negative (indefinite) curvature still yield a descent direction, and the + **relative** floor ``lam * max|diag|`` caps the preconditioner ratio at + ``1/lam`` (default 100:1). This is what keeps the seeded first step safe far + from the minimum: an absolute floor lets a near-zero-curvature (flat) + direction get an unbounded ``1/lam`` preconditioner, so the Newton-scaled + ``t=1`` step blows those directions to non-finite geometry. The relative + floor bounds the worst-conditioned direction's step to ``1/lam`` times the + stiffest direction's, keeping the whole step in a trust-region-like range. + + Parameters + ---------- + group_sizes : sequence of int, optional + Contiguous group lengths summing to ``diag.numel()`` (e.g. one entry per + ``nn.Parameter``). When given, the relative floor is computed **per group** + (each slice floored by its own ``max|diag|``) instead of globally. This + matters when groups have curvatures of very different magnitude -- e.g. + the few scaler params each touch all reflections and have far larger + loss-curvature than a single coordinate, so a global floor would crush the + body group to a near-zero step. Per-group flooring gives each group its own + Newton scale. + """ + if group_sizes is None: + return _precond_slice(diag, lam) + if int(sum(group_sizes)) != diag.numel(): + raise ValueError( + f"group_sizes sum {int(sum(group_sizes))} != diag numel {diag.numel()}" + ) + out = torch.empty_like(diag) + off = 0 + for n in group_sizes: + out[off:off + n] = _precond_slice(diag[off:off + n], lam) + off += n + return out + + +def _precond_slice(diag: torch.Tensor, lam: float) -> torch.Tensor: + if diag.numel() == 0: # empty group (e.g. absent aniso U / occupancy leaf) + return diag + a = diag.abs() + amax = a.max() + # abs backstop for the degenerate all-flat case (amax == 0) + floor = torch.clamp(lam * amax, min=torch.finfo(a.dtype).eps) + return 1.0 / torch.clamp(a, min=floor) + + +def hessian_diagonal_preconditioner( + aggregate: Callable[[], torch.Tensor], + params: Sequence[torch.nn.Parameter], + *, + n_probes: int = 16, + lam: float = 1e-2, + per_group: bool = False, + generator: Optional[torch.Generator] = None, + probe: str = "rademacher", +) -> torch.Tensor: + """Convenience: :func:`hessian_diagonal` then :func:`preconditioner_from_diagonal`. + + When ``per_group=True`` the relative floor is applied per ``params`` entry + (each ``nn.Parameter`` floored by its own ``max|diag|``) -- use this when the + params span groups of very different curvature magnitude (e.g. body + scaler). + """ + params = list(params) + diag = hessian_diagonal( + aggregate, + params, + n_probes=n_probes, + generator=generator, + probe=probe, + ) + group_sizes = [p.numel() for p in params] if per_group else None + return preconditioner_from_diagonal(diag, lam=lam, group_sizes=group_sizes) diff --git a/torchref/refinement/optimizers/seeded_lbfgs.py b/torchref/refinement/optimizers/seeded_lbfgs.py new file mode 100644 index 00000000..f05ca29d --- /dev/null +++ b/torchref/refinement/optimizers/seeded_lbfgs.py @@ -0,0 +1,300 @@ +"""SeededLBFGS: L-BFGS whose first step is a diagonal-Newton step. + +Stock :class:`torch.optim.LBFGS` initialises its inverse-Hessian approximation +with a *scalar* (``H_diag = ys / y.y``), so its very first search direction is +plain steepest descent ``d = -g``. That is badly conditioned when the parameter +blocks have very different curvature scales -- e.g. joint refinement of atomic +coordinates (Angstrom), ADPs (Angstrom^2 / log) and occupancies (sigmoid-logit) +in a single L-BFGS call. + +``SeededLBFGS`` accepts a per-parameter *inverse-curvature preconditioner* +``precond_i = 1 / (|H_ii| + lambda)`` (the diagonal of the loss Hessian, +regularised) and uses it to seed the **first** inner iteration only: + + d_0 = -(precond .* g) # diagonal-Newton direction + t_0 = lr # Newton scale -> start line search at t = 1 + +Every subsequent iteration is identical to stock L-BFGS (scalar ``H_diag`` seed, +strong-Wolfe line search, two-loop recursion). The seed fires exactly while the +curvature history is empty (``state["n_iter"] == 1``), which resets whenever +``optimizer.state.clear()`` is called -- so re-seeding a fresh block is just +``clear()`` + :meth:`set_init_hess_diag`. + +The diagonal itself is produced by +:mod:`torchref.refinement.optimizers.curvature` (Hutchinson estimate via analytic +double-backward Hessian-vector products under ``use_portable()``); this optimizer is +agnostic to how the preconditioner was obtained. + +Notes +----- +* With ``init_hess_diag=None`` the trajectory is **bit-identical** to + :class:`torch.optim.LBFGS` (both changed branches fall through to the exact base + expressions). +* With an all-ones preconditioner the first *direction* equals stock steepest + descent, but the first trial step length differs (``t = lr`` vs the + ``min(1, 1/||g||_1)`` heuristic), so the trajectory is *not* bit-identical -- + this is deliberate (the diagonal supplies Newton scale). +""" + +from __future__ import annotations + +from typing import Optional + +import torch +from torch.optim.lbfgs import _strong_wolfe, _to_scalar + + +class SeededLBFGS(torch.optim.LBFGS): + """L-BFGS that seeds its first step from the diagonal of the Hessian. + + Constructed exactly like :class:`torch.optim.LBFGS`. Before running, hand it + the inverse-curvature preconditioner via :meth:`set_init_hess_diag` (a flat + tensor in ``_gather_flat_grad`` order over the optimizer's parameters). Pass + ``None`` (the default) to behave exactly like stock L-BFGS. + """ + + def __init__(self, *args, init_hess_diag: Optional[torch.Tensor] = None, **kwargs): + super().__init__(*args, **kwargs) + self._init_hess_diag: Optional[torch.Tensor] = None + if init_hess_diag is not None: + self.set_init_hess_diag(init_hess_diag) + + def set_init_hess_diag(self, precond: Optional[torch.Tensor]) -> None: + """Set the flat inverse-curvature preconditioner ``1/(|diag|+lambda)``. + + Parameters + ---------- + precond : torch.Tensor or None + Flat tensor aligned to ``cat(p.view(-1) for p in params)`` (the same + layout as :meth:`_gather_flat_grad`). ``None`` disables seeding + (stock L-BFGS behaviour). The seed is applied to the first inner + iteration of the next :meth:`step` call whose curvature history is + empty. + """ + if precond is None: + self._init_hess_diag = None + return + if precond.dim() != 1: + precond = precond.reshape(-1) + expected = self._numel() + if precond.numel() != expected: + raise ValueError( + f"init_hess_diag numel {precond.numel()} != param numel {expected}" + ) + self._init_hess_diag = precond + + @torch.no_grad() + def step(self, closure): # type: ignore[override] + """Perform a single optimization step. + + Identical to :meth:`torch.optim.LBFGS.step` except that, on the first + inner iteration when a seed is set, the search direction is the + diagonal-Newton step ``-(precond .* g)`` and the initial trial step + length is ``lr``. + + Args: + closure (Callable): A closure that reevaluates the model and returns + the loss. + """ + assert len(self.param_groups) == 1 + + # Make sure the closure is always called with grad enabled + closure = torch.enable_grad()(closure) + + group = self.param_groups[0] + lr = _to_scalar(group["lr"]) + max_iter = group["max_iter"] + max_eval = group["max_eval"] + tolerance_grad = group["tolerance_grad"] + tolerance_change = group["tolerance_change"] + line_search_fn = group["line_search_fn"] + history_size = group["history_size"] + + # NOTE: LBFGS has only global state, but we register it as state for + # the first param, because this helps with casting in load_state_dict + state = self.state[self._params[0]] + state.setdefault("func_evals", 0) + state.setdefault("n_iter", 0) + + # evaluate initial f(x) and df/dx + orig_loss = closure() + loss = float(orig_loss) + current_evals = 1 + state["func_evals"] += 1 + + flat_grad = self._gather_flat_grad() + opt_cond = flat_grad.abs().max() <= tolerance_grad + + # optimal condition + if opt_cond: + return orig_loss + + # tensors cached in state (for tracing) + d = state.get("d") + t = state.get("t") + old_dirs = state.get("old_dirs") + old_stps = state.get("old_stps") + ro = state.get("ro") + H_diag = state.get("H_diag") + prev_flat_grad = state.get("prev_flat_grad") + prev_loss = state.get("prev_loss") + + # Resolve the diagonal seed once for this .step() call. Only consulted on + # the first inner iteration (empty history); cast to the grad's layout. + seed = self._init_hess_diag + if seed is not None: + seed = seed.to(device=flat_grad.device, dtype=flat_grad.dtype) + + n_iter = 0 + # optimize for a max of max_iter iterations + while n_iter < max_iter: + # keep track of nb of iterations + n_iter += 1 + state["n_iter"] += 1 + + ############################################################ + # compute gradient descent direction + ############################################################ + if state["n_iter"] == 1: + if seed is not None: + # ---- seeded first step: diagonal-Newton direction ---- + d = flat_grad.neg().mul_(seed) + else: + d = flat_grad.neg() + old_dirs = [] + old_stps = [] + ro = [] + H_diag = 1 + else: + # do lbfgs update (update memory) + y = flat_grad.sub(prev_flat_grad) + s = d.mul(t) + ys = y.dot(s) # y*s + if ys > 1e-10: + # updating memory + if len(old_dirs) == history_size: + # shift history by one (limited-memory) + old_dirs.pop(0) + old_stps.pop(0) + ro.pop(0) + + # store new direction/step + old_dirs.append(y) + old_stps.append(s) + ro.append(1.0 / ys) + + # update scale of initial Hessian approximation + H_diag = ys / y.dot(y) # (y*y) + + # compute the approximate (L-BFGS) inverse Hessian + # multiplied by the gradient + num_old = len(old_dirs) + + if "al" not in state: + state["al"] = [None] * history_size + al = state["al"] + + # iteration in L-BFGS loop collapsed to use just one buffer + q = flat_grad.neg() + for i in range(num_old - 1, -1, -1): + al[i] = old_stps[i].dot(q) * ro[i] + q.add_(old_dirs[i], alpha=-al[i]) + + # multiply by initial Hessian + # r/d is the final direction + d = r = torch.mul(q, H_diag) + for i in range(num_old): + be_i = old_dirs[i].dot(r) * ro[i] + r.add_(old_stps[i], alpha=al[i] - be_i) + + if prev_flat_grad is None: + prev_flat_grad = flat_grad.clone(memory_format=torch.contiguous_format) + else: + prev_flat_grad.copy_(flat_grad) + prev_loss = loss + + ############################################################ + # compute step length + ############################################################ + # reset initial guess for step size + if state["n_iter"] == 1: + if seed is not None: + # diagonal seed already carries Newton scale -> t = 1 (lr) + t = lr + else: + t = min(1.0, 1.0 / flat_grad.abs().sum()) * lr + else: + t = lr + + # directional derivative + gtd = flat_grad.dot(d) # g * d + + # directional derivative is below tolerance + if gtd > -tolerance_change: + break + + # optional line search: user function + ls_func_evals = 0 + if line_search_fn is not None: + # perform line search, using user function + if line_search_fn != "strong_wolfe": + raise RuntimeError("only 'strong_wolfe' is supported") + else: + x_init = self._clone_param() + + def obj_func(x, t, d): + return self._directional_evaluate(closure, x, t, d) + + loss, flat_grad, t, ls_func_evals = _strong_wolfe( + obj_func, x_init, t, d, loss, flat_grad, gtd + ) + self._add_grad(t, d) + opt_cond = flat_grad.abs().max() <= tolerance_grad + else: + # no line search, simply move with fixed-step + self._add_grad(t, d) + if n_iter != max_iter: + # re-evaluate function only if not in last iteration + # the reason we do this: in a stochastic setting, + # no use to re-evaluate that function here + with torch.enable_grad(): + loss = closure() + loss = float(loss) + flat_grad = self._gather_flat_grad() + opt_cond = flat_grad.abs().max() <= tolerance_grad + ls_func_evals = 1 + + # update func eval + current_evals += ls_func_evals + state["func_evals"] += ls_func_evals + + ############################################################ + # check conditions + ############################################################ + if n_iter == max_iter: + break + + if current_evals >= max_eval: + break + + # optimal condition + if opt_cond: + break + + # lack of progress + if d.mul(t).abs().max() <= tolerance_change: + break + + if abs(loss - prev_loss) < tolerance_change: + break + + state["d"] = d + state["t"] = t + state["old_dirs"] = old_dirs + state["old_stps"] = old_stps + state["ro"] = ro + state["H_diag"] = H_diag + state["prev_flat_grad"] = prev_flat_grad + state["prev_loss"] = prev_loss + + return orig_loss diff --git a/torchref/refinement/targets/adp/base.py b/torchref/refinement/targets/adp/base.py index 3e8940c6..68675a9a 100644 --- a/torchref/refinement/targets/adp/base.py +++ b/torchref/refinement/targets/adp/base.py @@ -39,6 +39,7 @@ def __init__( self, model: "Model" = None, verbose: int = 0, + device=None, **kwargs, ): - super().__init__(model, verbose) + super().__init__(model, verbose, device=device) diff --git a/torchref/refinement/targets/adp/locality.py b/torchref/refinement/targets/adp/locality.py index 9dbf41ad..b43ec182 100644 --- a/torchref/refinement/targets/adp/locality.py +++ b/torchref/refinement/targets/adp/locality.py @@ -60,18 +60,22 @@ def __init__( sigma_aniso: float = 0.5, exclude_bonded: bool = True, verbose: int = 0, + device=None, ): - super().__init__(model, verbose) - self.register_buffer( - "_k_neighbors", torch.tensor(k_neighbors, dtype=torch.int64) - ) - self.register_buffer("_correlation_length", torch.tensor(correlation_length)) - self.register_buffer("_scale", torch.tensor(scale)) + super().__init__(model, verbose, device=device) + # Host-side, deliberately not buffers: all three are only ever reached + # through the properties below, which immediately ``.item()`` them, and + # their consumers (k-NN sizing, the exp() falloff, stats reporting) are + # host-side. Keeping them on the device bought a sync per access. + self._k_neighbors = int(k_neighbors) + self._correlation_length = float(correlation_length) + self._scale = float(scale) # Sigma for the deviatoric (anisotropy) channel; only used when # anisotropic atoms are present. Dimensionless and on the same scale as # the magnitude channel's fixed 0.5 log-sigma (the channel restrains - # fractional anisotropy dev/B_eq, the analogue of log B_eq). - self.register_buffer("_sigma_aniso", torch.tensor(sigma_aniso)) + # fractional anisotropy dev/B_eq, the analogue of log B_eq). This one + # *is* a buffer: it is handed to adp_locality_aniso_math as a tensor. + self._register_scalar("_sigma_aniso", float(sigma_aniso)) self.exclude_bonded = exclude_bonded # Cache for neighbor indices and distances @@ -79,29 +83,48 @@ def __init__( self._neighbor_distances = None # (N, k_neighbors) self._last_xyz_hash = None + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + """Absorb the retired ``_k_neighbors`` / ``_correlation_length`` / + ``_scale`` buffers. + + All three are host-side Python scalars now (see ``__init__``). + Checkpoints predating that change still carry them, and a + ``strict=True`` load would reject them as unexpected keys -- so restore + their values rather than discarding a user's tuning. + """ + for legacy, cast in ( + ("_k_neighbors", int), + ("_correlation_length", float), + ("_scale", float), + ): + saved = state_dict.pop(prefix + legacy, None) + if saved is not None: + setattr(self, legacy, cast(saved.item())) + return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + @property def k_neighbors(self) -> int: - return self._k_neighbors.item() + return self._k_neighbors @k_neighbors.setter def k_neighbors(self, value: int): - self._k_neighbors.fill_(value) + self._k_neighbors = int(value) @property def correlation_length(self) -> float: - return self._correlation_length.item() + return self._correlation_length @correlation_length.setter def correlation_length(self, value: float): - self._correlation_length.fill_(value) + self._correlation_length = float(value) @property def scale(self) -> float: - return self._scale.item() + return self._scale @scale.setter def scale(self, value: float): - self._scale.fill_(value) + self._scale = float(value) @property def sigma_aniso(self) -> float: @@ -287,11 +310,6 @@ def forward(self, recompute_neighbors: bool = False) -> torch.Tensor: # deviatoric/anisotropy channel). Isotropic-only models keep the # original B-factor path (numerically unchanged, zero overhead). if not getattr(self.model, "_aniso_is_empty", True): - if (self._sigma_aniso.device != device - or self._sigma_aniso.dtype != adp.dtype): - self._sigma_aniso = self._sigma_aniso.to( - device=device, dtype=adp.dtype - ) u6 = self.model.adp_u6() return adp_locality_aniso_math( u6, indices, distances, self._sigma_aniso diff --git a/torchref/refinement/targets/adp/similarity.py b/torchref/refinement/targets/adp/similarity.py index 78335a1c..45dd914b 100644 --- a/torchref/refinement/targets/adp/similarity.py +++ b/torchref/refinement/targets/adp/similarity.py @@ -42,12 +42,15 @@ class ADPSimilarityTarget(ADPTarget): def __init__( self, model: "Model" = None, simu_sigma: float = 2.0, - simu_sigma_aniso: float = 1.0, verbose: int = 0 + simu_sigma_aniso: float = 1.0, verbose: int = 0, device=None ): - super().__init__(model, verbose) - # Register simu-specific sigma as buffer (separate from base sigma) - self.register_buffer("_simu_sigma", torch.tensor(simu_sigma)) - self.register_buffer("_simu_sigma_aniso", torch.tensor(simu_sigma_aniso)) + super().__init__(model, verbose, device=device) + # Both sigmas are handed to adp_simu_math / adp_simu_aniso_math, which + # dispatch to Triton on CUDA float32. Allocating them on the target's + # resolved device and dtype is what lets the lazy repair inside + # forward() go away. + self._register_scalar("_simu_sigma", float(simu_sigma)) + self._register_scalar("_simu_sigma_aniso", float(simu_sigma_aniso)) @property def simu_sigma(self) -> float: @@ -96,25 +99,10 @@ def forward(self) -> torch.Tensor: adp_t = self.model.adp() if pair_indices.shape[0] == 0: return torch.zeros((), device=adp_t.device, dtype=adp_t.dtype) - # Lazily move the ``_simu_sigma`` buffer onto the model's device - # the first time we reach here. Once moved, subsequent forwards - # (and CUDA-Graph captures) skip the device transfer — calling - # ``.to()`` on a CPU buffer inside a capture region triggers a - # ``cudaErrorStreamCaptureUnsupported``. - if (self._simu_sigma.device != adp_t.device - or self._simu_sigma.dtype != adp_t.dtype): - self._simu_sigma = self._simu_sigma.to( - device=adp_t.device, dtype=adp_t.dtype, - ) # Anisotropic atoms present: restrain the full U tensors. Isotropic-only # models keep the original (Triton-accelerated) B-factor path so they # pay nothing and are numerically unchanged. if not getattr(self.model, "_aniso_is_empty", True): - if (self._simu_sigma_aniso.device != adp_t.device - or self._simu_sigma_aniso.dtype != adp_t.dtype): - self._simu_sigma_aniso = self._simu_sigma_aniso.to( - device=adp_t.device, dtype=adp_t.dtype, - ) u6 = self.model.adp_u6() return adp_simu_aniso_math( u6, pair_indices, self._simu_sigma, self._simu_sigma_aniso diff --git a/torchref/refinement/targets/base.py b/torchref/refinement/targets/base.py index 237b53af..7ea97cd2 100644 --- a/torchref/refinement/targets/base.py +++ b/torchref/refinement/targets/base.py @@ -21,7 +21,9 @@ from torch import nn from torch.special import i0 +from torchref.config import get_float_dtype, normalize_device from torchref.utils.device_mixin import DeviceMixin +from torchref.utils.device_resolution import resolve_device from torchref.utils.stats import ( VERBOSITY_DEBUG, VERBOSITY_DETAILED, @@ -80,6 +82,7 @@ class Target(DeviceMixin, nn.Module): def __init__( self, verbose: int = 0, + device=None, **kwargs, ): """ @@ -89,9 +92,79 @@ def __init__( ---------- verbose : int, optional Verbosity level. Default is 0. + device : torch.device, optional + Where this target allocates. Defaults to the configured default; + :meth:`_adopt_device` refines it from whatever model / data / + scaler the subclass is given. """ super().__init__() self.verbose = verbose + # Seeded here for two distinct reasons. + # + # (1) Subclasses allocate their tunables in ``__init__``, so + # ``device=self.device`` / ``dtype=self.dtype_float`` must already + # have an answer at that point. Before this existed, every scalar + # tunable in the package landed on CPU float32 regardless of the + # model's device or the configured dtype. + # + # (2) ``DeviceMixin._refresh_device_trackers`` early-returns unless one + # of these names is already in ``obj.__dict__``. Merely defining + # them is what switches the mixin on for targets -- it was + # previously a complete no-op for every target in the package. + self.device = normalize_device(device) + self.dtype_float = get_float_dtype() + + # ---- device / dtype resolution -------------------------------------- + + def _adopt_device(self, *sources, device=None): + """Reconcile this target with the device-bearing objects it wraps. + + Call *after* the sources are attached and *before* allocating any + buffer. ``None`` sources are dropped, so the empty-init path used by + ``load_state_dict`` keeps the seeded default, and the method stays + re-runnable if a source is attached later. + + ``self`` is deliberately **not** passed to + :func:`~torchref.utils.resolve_device` -- unlike + ``ScalerBase.set_data``, which does pass it. A freshly constructed + target owns no tensors, so there is nothing of its own to reconcile, + and a target's state is a handful of scalars against a ``Model``'s + entire structure. Letting an empty shell's default device outvote a + model already on the accelerator would move megabytes to satisfy five + floats. + + Returns + ------- + torch.device + The resolved device (also stored on ``self.device``). + """ + present = [s for s in sources if s is not None] + if device is None and not present: + return self.device + self.device = resolve_device(*present, device=device) + for src in present: + src_dtype = getattr(src, "dtype_float", None) + if isinstance(src_dtype, torch.dtype): + self.dtype_float = src_dtype + break + return self.device + + def _register_scalar(self, name: str, value, dtype=None) -> None: + """Register a 0-dim tunable on this target's device and float dtype. + + Scalar tunables reach Triton kernels as tensors, where a CPU-resident + buffer is a raw pointer into host memory rather than a promotable + scalar. Allocating them correctly here is what lets the lazy ``.to()`` + repairs inside ``forward()`` go away. + """ + self.register_buffer( + name, + torch.tensor( + value, + device=self.device, + dtype=self.dtype_float if dtype is None else dtype, + ), + ) def forward(self) -> torch.Tensor: """Compute and return the loss. Override in subclasses.""" @@ -181,6 +254,7 @@ def __init__( self, model: "Model" = None, verbose: int = 0, + device=None, **kwargs, ): """ @@ -192,11 +266,16 @@ def __init__( Reference to the Model object (optional for empty init). verbose : int, optional Verbosity level. Default is 0. + device : torch.device, optional + Explicit device. When omitted, follows ``model``. """ - super().__init__(verbose=verbose) + super().__init__(verbose=verbose, device=device) # Register model as a proper submodule (not in state_dict but handles device) # Use add_module to allow None values self.add_module("_model", model) + # Follow the model rather than the global default, so subclass buffers + # allocated below this line land beside the coordinates they act on. + self._adopt_device(model, device=device) @property def model(self) -> "Model": @@ -268,6 +347,7 @@ def __init__( model: "Model" = None, scaler: "Scaler" = None, verbose: int = 0, + device=None, **kwargs, ): """ @@ -284,12 +364,21 @@ def __init__( Reference to the Scaler object. verbose : int, optional Verbosity level. Default is 0. + device : torch.device, optional + Explicit device. When omitted, model / data / scaler are + reconciled onto one device, the model winning on disagreement. """ - super().__init__(verbose=verbose) - # Register as proper submodules (allows None values) + super().__init__(verbose=verbose, device=device) + # Register as proper submodules (allows None values). Note ``_data`` is + # a plain attribute: ReflectionData is a dataclass, not an nn.Module, + # so add_module would reject it -- but DeviceMixin still walks it. self.add_module("_model", model) self._data = data self.add_module("_scaler", scaler) + # The forward path mixes tensors from all three, so pin them together + # here. Order is precedence: ``resolve_device`` is first-wins, and the + # model is what the rest of the pipeline follows. + self._adopt_device(model, data, scaler, device=device) @property def model(self) -> "Model": diff --git a/torchref/refinement/targets/geometry/base.py b/torchref/refinement/targets/geometry/base.py index d9908591..04c7c794 100644 --- a/torchref/refinement/targets/geometry/base.py +++ b/torchref/refinement/targets/geometry/base.py @@ -39,9 +39,10 @@ def __init__( self, model: "Model" = None, verbose: int = 0, + device=None, **kwargs, ): - super().__init__(model, verbose) + super().__init__(model, verbose, device=device) def stats(self) -> Dict[str, StatEntry]: """ diff --git a/torchref/refinement/targets/geometry/non_bonded.py b/torchref/refinement/targets/geometry/non_bonded.py index d4f457bb..1d61abc5 100644 --- a/torchref/refinement/targets/geometry/non_bonded.py +++ b/torchref/refinement/targets/geometry/non_bonded.py @@ -109,6 +109,7 @@ def __init__( rebuild_threshold: float = 1.0, verbose: int = 0, scale: float = 10.0, + device=None, ): """ Initialize non-bonded target. @@ -142,23 +143,50 @@ def __init__( Public attribute stored as ``self.scale``. Default is 10.0. Set but not consumed by ``forward()`` (informational only). """ - super().__init__(model, verbose) + super().__init__(model, verbose, device=device) self.mode = mode self.scale = scale - # Register sigma / r_exp / buffer as buffers so .to(device) moves them. - self.register_buffer("_sigma_vdw", torch.tensor(float(sigma))) - self.register_buffer("_r_exp", torch.tensor(float(r_exp))) - self.register_buffer("_buffer", torch.tensor(float(buffer))) - self.register_buffer( - "_rebuild_threshold", torch.tensor(float(rebuild_threshold)) - ) + # Tunables that reach the kernel live on the target's resolved device + # and float dtype: the prolsq branch hands these straight to a Triton + # kernel, where a CPU tensor is a raw pointer into host memory rather + # than a promotable scalar. + self._register_scalar("_sigma_vdw", float(sigma)) + self._register_scalar("_r_exp", float(r_exp)) # c_rep: back-door override; by default derived from sigma so that # PROLSQ shape term equals v^p / (p * sigma^p). if c_rep is None: c_rep_val = 1.0 / (float(r_exp) * float(sigma) ** float(r_exp)) else: c_rep_val = float(c_rep) - self.register_buffer("_c_rep", torch.tensor(c_rep_val)) + self._register_scalar("_c_rep", c_rep_val) + # Host-side, deliberately not buffers. + # + # ``buffer`` is consumed as a Python float -- forward() already called + # ``float(self._buffer)`` before handing it to the kernel, where it + # becomes a compile-time constant. Storing it on the device only bought + # a device->host sync on the hot path. + # + # ``rebuild_threshold`` is only ever compared against an already + # ``.item()``-ed displacement in maintenance(). + self._buffer = float(buffer) + self._rebuild_threshold = float(rebuild_threshold) + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + """Absorb the retired ``_buffer`` / ``_rebuild_threshold`` buffers. + + Both are host-side Python floats now (see ``__init__``). Checkpoints + written before that change still carry them, and a ``strict=True`` load + would reject them as unexpected keys -- so restore their values rather + than discarding a user's tuning. + """ + for legacy, attr in ( + ("_buffer", "_buffer"), + ("_rebuild_threshold", "_rebuild_threshold"), + ): + saved = state_dict.pop(prefix + legacy, None) + if saved is not None: + setattr(self, attr, float(saved.item())) + return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) @property def c_rep(self) -> float: @@ -199,13 +227,13 @@ def r_exp(self, value: float): @property def buffer(self) -> float: - """Get distance buffer.""" - return self._buffer.item() + """Get distance buffer (host-side; see ``__init__``).""" + return self._buffer @buffer.setter def buffer(self, value: float): """Set distance buffer.""" - self._buffer.fill_(value) + self._buffer = float(value) def maintenance(self) -> None: """Rebuild the VDW pair list if any ASU atom drifted too far. @@ -245,12 +273,12 @@ def maintenance(self) -> None: max_disp_sq = (delta * delta).sum(dim=-1).max() thresh_sq = self._rebuild_threshold * self._rebuild_threshold - if max_disp_sq.item() <= thresh_sq.item(): + if max_disp_sq.item() <= thresh_sq: return # within slack — nothing to do if self.verbose > 0: max_disp = float(max_disp_sq.item()) ** 0.5 - thresh = float(self._rebuild_threshold.item()) + thresh = self._rebuild_threshold print( f" VDW rebuild: max drift {max_disp:.2f} Å > " f"threshold {thresh:.2f} Å" @@ -358,7 +386,7 @@ def forward(self) -> torch.Tensor: self.model.cell.fractional_matrix, self.model.cell.inv_fractional_matrix, self._c_rep, self._r_exp, - float(self._buffer), self._sigma_vdw, + self._buffer, self._sigma_vdw, ) # Compute positions (handles symmetry transparently) diff --git a/torchref/refinement/targets/similarity.py b/torchref/refinement/targets/similarity.py index 01872a92..83edc2c6 100644 --- a/torchref/refinement/targets/similarity.py +++ b/torchref/refinement/targets/similarity.py @@ -72,11 +72,26 @@ def __init__( model_light: "Model" = None, alpha: float = 2.0, verbose: int = 0, + device=None, ): - super().__init__(verbose=verbose) + super().__init__(verbose=verbose, device=device) self.add_module("_model_dark", model_dark) self.add_module("_model_light", model_light) - self.register_buffer("_alpha", torch.tensor(alpha)) + self._adopt_device(model_dark, model_light, device=device) + self._register_scalar("_alpha", float(alpha)) + # Registered unconditionally, before the map is built. They used to be + # created only inside ``_build_atom_map``, which runs only when both + # models are present -- so on the empty-init path (the one that exists + # for ``load_state_dict``) they never existed at all: ``forward()`` + # raised AttributeError and a strict load of a populated checkpoint + # failed on unexpected keys. ``_build_atom_map`` now overwrites rather + # than creates. + self.register_buffer( + "_idx_dark", torch.zeros(0, dtype=torch.long, device=self.device) + ) + self.register_buffer( + "_idx_light", torch.zeros(0, dtype=torch.long, device=self.device) + ) if model_dark is not None and model_light is not None: self._build_atom_map() @@ -148,10 +163,10 @@ def _build_atom_map(self): "dark and light models" ) self.register_buffer( - "_idx_dark", torch.zeros(0, dtype=torch.long) + "_idx_dark", torch.zeros(0, dtype=torch.long, device=self.device) ) self.register_buffer( - "_idx_light", torch.zeros(0, dtype=torch.long) + "_idx_light", torch.zeros(0, dtype=torch.long, device=self.device) ) return @@ -168,13 +183,21 @@ def _build_atom_map(self): f"(dark={n_dark}, light={n_light})" ) + # On ``self.device``: these are 1-D index tensors, so unlike the 0-dim + # scalars they are not covered by PyTorch's scalar-promotion rule. A + # CPU-resident index against accelerator coordinates costs a host sync + # on every forward. self.register_buffer( "_idx_dark", - torch.tensor(merged["_idx_dark"].values, dtype=torch.long), + torch.tensor( + merged["_idx_dark"].values, dtype=torch.long, device=self.device + ), ) self.register_buffer( "_idx_light", - torch.tensor(merged["_idx_light"].values, dtype=torch.long), + torch.tensor( + merged["_idx_light"].values, dtype=torch.long, device=self.device + ), ) def forward(self) -> torch.Tensor: @@ -186,8 +209,9 @@ def forward(self) -> torch.Tensor: Scalar mean loss over all matched atom pairs. """ if len(self._idx_dark) == 0: - device = self._alpha.device - return torch.tensor(0.0, device=device) + # ``self.device``, not ``self._alpha.device``: an empty target must + # still hand back a loss on the refinement's device. + return torch.tensor(0.0, device=self.device, dtype=self.dtype_float) xyz_dark = self._model_dark.xyz() xyz_light = self._model_light.xyz() diff --git a/torchref/refinement/targets/xray/base.py b/torchref/refinement/targets/xray/base.py index c0ccb425..b74f606e 100644 --- a/torchref/refinement/targets/xray/base.py +++ b/torchref/refinement/targets/xray/base.py @@ -67,6 +67,7 @@ def __init__( use_work_set: bool = True, verbose: int = 0, use_set: str = None, + device=None, ): """ Initialize X-ray target. @@ -89,7 +90,9 @@ def __init__( verbose : int, optional Verbosity level. Default is 0. """ - super().__init__(data=data, model=model, scaler=scaler, verbose=verbose) + super().__init__( + data=data, model=model, scaler=scaler, verbose=verbose, device=device + ) # ``use_set`` (3-way: "work"/"free"/"val") is the canonical subset # selector; the legacy ``use_work_set`` bool maps onto it. When # ``use_set`` is not given explicitly, fall back to the bool so older diff --git a/torchref/refinement/targets/xray/bhattacharyya.py b/torchref/refinement/targets/xray/bhattacharyya.py index 04f3aeb9..fe3cb702 100644 --- a/torchref/refinement/targets/xray/bhattacharyya.py +++ b/torchref/refinement/targets/xray/bhattacharyya.py @@ -107,30 +107,49 @@ def __init__( verbose=verbose, use_set=use_set, ) - # log-spaced B grid + # log-spaced B grid. ``log_b_grid`` itself is a construction-time + # local, not a buffer: it is read twice below for the min/max and never + # touched again, so registering it only created something for ``.to()`` + # to carry around. log_b_grid = torch.linspace( - math.log(b_grid_min), math.log(b_grid_max), b_grid_n + math.log(b_grid_min), + math.log(b_grid_max), + b_grid_n, + device=self.device, + dtype=self.dtype_float, ) - self.register_buffer("log_b_grid", log_b_grid) self.register_buffer("b_grid", torch.exp(log_b_grid)) self._log_b_min = float(log_b_grid[0].item()) self._log_b_max = float(log_b_grid[-1].item()) self._log_b_step = (self._log_b_max - self._log_b_min) / (b_grid_n - 1) # Global σ_m scale (tunable, non-learnable) + self._register_scalar("sigma_m_scale", float(sigma_m_scale)) + # Populated by _initialize_cache() on first forward(). Empty, but still + # allocated on this target's device: a 0-element tensor reports a + # device, and a CPU one here reads as a split object. + empty = dict(device=self.device, dtype=self.dtype_float) + self.register_buffer("exp_table", torch.empty(0, **empty)) # (b_grid_n, N_refl) + self.register_buffer("s_sq_per_refl", torch.empty(0, **empty)) # (N_refl,) + self.register_buffer("s_4_per_refl", torch.empty(0, **empty)) # (N_refl,) + self.register_buffer("f_sq_kh", torch.empty(0, **empty)) # (K, N_refl) + self.register_buffer("g_w_table", torch.empty(0, **empty)) # (K, b_grid_n) + self.register_buffer("g_4_table", torch.empty(0, **empty)) # (K, b_grid_n) self.register_buffer( - "sigma_m_scale", torch.tensor(float(sigma_m_scale)) + "atom_to_element", torch.empty(0, dtype=torch.long, device=self.device) ) - # Populated by _initialize_cache() on first forward() - self.register_buffer("exp_table", torch.empty(0)) # (b_grid_n, N_refl) - self.register_buffer("s_sq_per_refl", torch.empty(0)) # (N_refl,) - self.register_buffer("s_4_per_refl", torch.empty(0)) # (N_refl,) - self.register_buffer("f_sq_kh", torch.empty(0)) # (K, N_refl) - self.register_buffer("g_w_table", torch.empty(0)) # (K, b_grid_n) - self.register_buffer("g_4_table", torch.empty(0)) # (K, b_grid_n) - self.register_buffer("atom_to_element", torch.empty(0, dtype=torch.long)) - self.register_buffer("sigma_d_mean", torch.tensor(0.0)) + self._register_scalar("sigma_d_mean", 0.0) self._initialized = False + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + """Drop the retired ``log_b_grid`` buffer key. + + It is a construction-time local now (derived from the same + ``b_grid_min``/``max``/``n`` the constructor already has), so a + ``strict=True`` load of an older checkpoint would reject it. + """ + state_dict.pop(prefix + "log_b_grid", None) + return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + # ------------------------------------------------------------------ # Cache initialisation (called once, on first forward) # ------------------------------------------------------------------ diff --git a/torchref/refinement/targets/xray/least_squares.py b/torchref/refinement/targets/xray/least_squares.py index 4dbe340f..5fdff64e 100644 --- a/torchref/refinement/targets/xray/least_squares.py +++ b/torchref/refinement/targets/xray/least_squares.py @@ -46,6 +46,7 @@ def __init__( n_bins: int = 20, verbose: int = 0, use_set: str = None, + device=None, ): if scale_mode not in ("scaler", "binwise_optimal"): raise ValueError( @@ -64,6 +65,7 @@ def __init__( use_work_set=use_work_set, verbose=verbose, use_set=use_set, + device=device, ) self.weighting = weighting self.scale_mode = scale_mode diff --git a/torchref/refinement/targets/xray/maximum_likelihood.py b/torchref/refinement/targets/xray/maximum_likelihood.py index af36f00c..65878db7 100644 --- a/torchref/refinement/targets/xray/maximum_likelihood.py +++ b/torchref/refinement/targets/xray/maximum_likelihood.py @@ -8,7 +8,6 @@ epsilon_from_hkl, ml_xray_loss_beta_math, ) -from torchref.utils.device_resolution import resolve_device from .base import XrayTarget from .gaussian import GaussianXrayTarget @@ -178,10 +177,10 @@ def create_xray_target( ) mode = "ml" - # Pin model/data/scaler onto one device before constructing the - # target — its forward path mixes tensors from all three. - resolve_device(model, data, scaler, device=device) - + # Device reconciliation is ``DataTarget.__init__``'s job now (it calls + # ``_adopt_device(model, data, scaler)`` before allocating anything), so + # doing it again here would be a second copy of the same policy, free to + # drift. ``device`` is forwarded instead. kwargs = dict( data=data, model=model, @@ -189,6 +188,7 @@ def create_xray_target( use_work_set=use_work_set, verbose=verbose, use_set=use_set, + device=device, ) if mode == "gaussian": return GaussianXrayTarget(**kwargs) diff --git a/torchref/restraints/hydrogen_topology.py b/torchref/restraints/hydrogen_topology.py index b4b8e94f..ecd6f0d2 100644 --- a/torchref/restraints/hydrogen_topology.py +++ b/torchref/restraints/hydrogen_topology.py @@ -17,7 +17,8 @@ import torch from torch import nn -from torchref.config import dtypes, get_default_device +from torchref.config import dtypes, normalize_device +from torchref.utils.device_resolution import resolve_device from torchref.utils.device_mixin import DeviceMixin # --------------------------------------------------------------------------- @@ -83,9 +84,14 @@ class HydrogenTopology(DeviceMixin, nn.Module): ``build_h_candidate_pairs`` and checked via :attr:`has_candidates`. """ - def __init__(self): + def __init__(self, device=None): super().__init__() - # Buffers are registered by build_hydrogen_topology() + # Buffers are registered later by build_hydrogen_topology(), so this + # object is tensor-free at construction. The tracker still has to exist: + # ``DeviceMixin._refresh_device_trackers`` only maintains attributes + # already present in ``__dict__``, and callers reconcile against + # ``h_topo.device`` before attaching buffers to it. + self.device = normalize_device(device) @property def n_hydrogens(self) -> int: @@ -250,8 +256,7 @@ def build_hydrogen_topology( HydrogenTopology Module with registered buffer tensors. """ - if device is None: - device = get_default_device() + device = normalize_device(device) cache = _load_cif_hydrogen_info(pdb, verbose) model_names = pdb["name"].astype(str).str.strip().values @@ -366,7 +371,9 @@ def build_hydrogen_topology( acc_chainid_enc.append(chainid_enc) acc_resseq.append(resseq) - topo = HydrogenTopology() + # Seed the tracker with the device its buffers are about to be built on, + # so a later ``resolve_device(h_topo, ...)`` sees the truth. + topo = HydrogenTopology(device=device) n_h_total = len(acc_parent_idx) fdtype = dtypes.float @@ -631,6 +638,10 @@ def place_riding_hydrogens( ------- xyz_h : (N_h, 3) float tensor, differentiable w.r.t. xyz_heavy """ + # Function-local, matching every other target dispatch site: importing the gate at + # module scope would pull ``torchref.base.targets`` into ``torchref.restraints``. + from torchref.base.targets._dispatch import use_triton + N_h = topo.h_parent_idx.shape[0] if N_h == 0: return torch.zeros(0, 3, dtype=xyz_heavy.dtype, device=xyz_heavy.device) @@ -651,7 +662,13 @@ def place_riding_hydrogens( # backward (the H-VDW backward is a significant fraction of the # non-bonded fwd+bw cost). Falls back to the JIT-scripted eager # helper otherwise. - if xyz_heavy.is_cuda and xyz_heavy.dtype == torch.float32: + # + # Gated by the shared ``use_triton`` rather than a hand-rolled + # ``is_cuda and dtype == float32``. The inline check ignored the shared gate entirely, so + # a block that pinned the portable path still ran the Triton kernel here -- and since + # EAGER is the documented double-differentiable route, a Hessian taken through hydrogen + # placement was silently going through a first-order-only kernel. + if use_triton(xyz_heavy): try: from torchref.base.targets.triton.place_hydrogens import ( place_riding_hydrogens_triton, @@ -719,8 +736,10 @@ def build_h_candidate_pairs( device : torch.device verbose : int """ - if device is None: - device = get_default_device() + # Buffers are registered onto ``h_topo`` below, so follow it rather than + # the global default -- otherwise they attach to a module living somewhere + # else. + device = resolve_device(h_topo, device=device) n_h = h_topo.n_hydrogens n_heavy = len(pdb) diff --git a/torchref/restraints/restraints.py b/torchref/restraints/restraints.py index ada6ef82..f3da1fa7 100644 --- a/torchref/restraints/restraints.py +++ b/torchref/restraints/restraints.py @@ -29,7 +29,7 @@ read_cif, read_link_definitions, ) -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_float_dtype from torchref.utils.debug_utils import DebugMixin from torchref.utils.utils import TensorDict from torchref.utils.device_mixin import DeviceMixin diff --git a/torchref/scaling/collection_scaler.py b/torchref/scaling/collection_scaler.py index 2683ea4b..7b8a82e6 100644 --- a/torchref/scaling/collection_scaler.py +++ b/torchref/scaling/collection_scaler.py @@ -26,7 +26,7 @@ epsilon_from_hkl, ml_xray_loss_beta_math, ) -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_float_dtype from torchref.scaling.scaler_base import ScalerBase from torchref.scaling.solvent import SolventModel from torchref.utils.utils import ModuleReference @@ -80,10 +80,12 @@ def __init__( verbose: int = 1, device: torch.device = None, ): - if device is None: - device = get_default_device() # Bind to the dark/reference dataset for bins and scattering vectors dark_data = dataset_collection[model_collection.dark_key] + # Forward ``device`` through rather than resolving the global default + # here: a resolved default arrives at ``ScalerBase`` as an *explicit* + # device, which then drags ``dark_data`` onto it. Passing ``None`` + # lets ScalerBase derive from the data, which is the point. super().__init__( data=dark_data, nbins=nbins, diff --git a/torchref/scaling/scaler.py b/torchref/scaling/scaler.py index 59a8b3b9..d29ff0f4 100644 --- a/torchref/scaling/scaler.py +++ b/torchref/scaling/scaler.py @@ -19,7 +19,6 @@ import torch import torch.nn as nn -from torchref.config import get_default_device from torchref.io import ReflectionData from torchref.base.reciprocal import get_scattering_vectors from torchref.scaling.scaler_base import ScalerBase @@ -157,7 +156,15 @@ def set_model_and_data(self, model: "Model", data: ReflectionData): Model object for structure factor calculation. data : ReflectionData ReflectionData object with observed data. + + Notes + ----- + Receiver wins: the scaler already owns buffers by this point, so + ``model`` and ``data`` are reconciled onto *its* device. This is the + state-dict restore path, where the three objects are built separately + and can easily disagree. """ + resolve_device(self, model, data) # Set _model_ref directly: nn.Module.__setattr__ intercepts assignments # of nn.Module instances (like `self.model = model`) and registers them # as submodules, bypassing the property setter entirely. diff --git a/torchref/scaling/scaler_base.py b/torchref/scaling/scaler_base.py index 18706063..b94f967e 100644 --- a/torchref/scaling/scaler_base.py +++ b/torchref/scaling/scaler_base.py @@ -30,7 +30,7 @@ epsilon_from_hkl, ml_xray_loss_beta_math, ) -from torchref.config import get_complex_dtype, get_default_device, get_float_dtype +from torchref.config import get_complex_dtype, get_float_dtype from torchref.utils.autograd_ops import gather_with_index_add from torchref.utils.debug_utils import DebugMixin from torchref.utils.device_mixin import DeviceMixin @@ -158,7 +158,14 @@ def set_data(self, data: "ReflectionData"): ---------- data : ReflectionData ReflectionData object with observed data. + + Notes + ----- + Receiver wins: this scaler may already hold buffers, so ``data`` is + moved onto the scaler's device rather than the buffers being registered + from whatever device ``data`` happened to be on. """ + self.device = resolve_device(self, data) self._data = ModuleReference(data) if data.cell is not None: self.cell = data.cell diff --git a/torchref/scaling/solvent.py b/torchref/scaling/solvent.py index 131d036f..90a6c65c 100644 --- a/torchref/scaling/solvent.py +++ b/torchref/scaling/solvent.py @@ -11,9 +11,10 @@ ifft, ) from torchref.base.electron_density.main import _get_radius_offsets -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_float_dtype from torchref.utils.debug_utils import DebugMixin from torchref.utils.device_mixin import DeviceMixin +from torchref.utils.device_resolution import resolve_device from torchref.utils.utils import ModuleReference, TensorDict @@ -113,8 +114,11 @@ def __init__( super(SolventModel, self).__init__() if float_type is None: float_type = get_float_dtype() - if device is None: - device = get_default_device() + # Follow the model when no device is given. ``realspace.py`` constructs + # a SolventModel with only a model, so falling back to the global + # default here put the solvent grids on a different device than the + # structure they are computed from. + device = resolve_device(model, device=device) self.device = device self.verbose = verbose self.float_type = float_type diff --git a/torchref/symmetry/cell.py b/torchref/symmetry/cell.py index 18d42f0c..88f5104f 100644 --- a/torchref/symmetry/cell.py +++ b/torchref/symmetry/cell.py @@ -14,8 +14,8 @@ from torchref.config import ( NYQUIST_OVERSAMPLING, - get_default_device, get_float_dtype, + normalize_device, ) from torchref.utils.device_mixin import _NonModuleDeviceMixin @@ -76,8 +76,7 @@ def __init__( """ if dtype is None: dtype = get_float_dtype() - if device is None: - device = get_default_device() + device = normalize_device(device) # Convert to tensor first to get shape if isinstance(data, torch.Tensor): tensor = data.to(dtype=dtype, device=device) diff --git a/torchref/symmetry/map_symmetry.py b/torchref/symmetry/map_symmetry.py index ae5dc613..457641e7 100644 --- a/torchref/symmetry/map_symmetry.py +++ b/torchref/symmetry/map_symmetry.py @@ -14,7 +14,7 @@ import torch import torch.nn as nn -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_float_dtype, normalize_device from torchref.symmetry.spacegroup import SpaceGroup, SpaceGroupLike from torchref.utils.device_mixin import DeviceMixin @@ -57,8 +57,12 @@ def MapSymmetry( """ if dtype_float is None: dtype_float = get_float_dtype() - if device is None: - device = get_default_device() + # ``cell_params`` is documented as a tensor, so follow it when no device is + # given rather than jumping to the global default and leaving the caller's + # cell behind. + if device is None and isinstance(cell_params, torch.Tensor): + device = cell_params.device + device = normalize_device(device) # Check grid compatibility symmetry = SpaceGroup(space_group, dtype=dtype_float, device=device) compat = symmetry.check_grid_compatibility(map_shape) @@ -117,14 +121,24 @@ def __init__( super().__init__() if dtype_float is None: dtype_float = get_float_dtype() - if device is None: - device = get_default_device() + if device is None and isinstance(cell_params, torch.Tensor): + device = cell_params.device self.dtype_float = dtype_float self.space_group = space_group self.map_shape = tuple(map_shape) - self.cell_params = cell_params self.verbose = verbose - self.device = device + self.device = normalize_device(device) + # Store ``cell_params`` *on* this module's device. It used to be kept + # exactly as handed in -- a plain attribute, so ``DeviceMixin`` could + # neither see it nor repair it, and ``self.device`` could disagree with + # it indefinitely. + if isinstance(cell_params, torch.Tensor): + cell_params = cell_params.to(device=self.device, dtype=self.dtype_float) + else: + cell_params = torch.as_tensor( + cell_params, device=self.device, dtype=self.dtype_float + ) + self.cell_params = cell_params self.symmetry = SpaceGroup( space_group, dtype=self.dtype_float, device=self.device diff --git a/torchref/symmetry/map_symmetry_interpolation.py b/torchref/symmetry/map_symmetry_interpolation.py index 3debb3a2..b11213a7 100644 --- a/torchref/symmetry/map_symmetry_interpolation.py +++ b/torchref/symmetry/map_symmetry_interpolation.py @@ -10,7 +10,7 @@ import torch.nn as nn import torch.nn.functional as F -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_float_dtype, normalize_device from torchref.symmetry.spacegroup import SpaceGroup from torchref.utils.device_mixin import DeviceMixin @@ -82,8 +82,7 @@ def __init__( super().__init__() if dtype_float is None: dtype_float = get_float_dtype() - if device is None: - device = get_default_device() + device = normalize_device(device) self.dtype_float = dtype_float self.space_group = space_group self.map_shape = tuple(map_shape) diff --git a/torchref/symmetry/reciprocal_symmetry.py b/torchref/symmetry/reciprocal_symmetry.py index ad7070b6..527fe7a8 100644 --- a/torchref/symmetry/reciprocal_symmetry.py +++ b/torchref/symmetry/reciprocal_symmetry.py @@ -30,7 +30,7 @@ import torch import torch.nn as nn -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_float_dtype, normalize_device from torchref.symmetry.spacegroup import SpaceGroup, SpaceGroupLike from torchref.utils.device_mixin import DeviceMixin @@ -133,8 +133,7 @@ def __init__( super().__init__() if dtype_float is None: dtype_float = get_float_dtype() - if device is None: - device = get_default_device() + device = normalize_device(device) self.dtype_float = dtype_float self.space_group = space_group self.grid_shape = tuple(grid_shape) @@ -772,8 +771,23 @@ def expand_hkl( hkl_transformed = torch.round(torch.matmul(hkl_float, recip_matrices[i].T)).to( torch.int32 ) - # Phase shift from translation: 2π × h·t - phase_shift = 2.0 * np.pi * torch.matmul(hkl_float, translations[i]) + # Phase shift from translation: -2π × h·t + # + # Derivation (do not "simplify" the sign away -- it was +2π here until + # 2026-07 and the error is invisible in P21/P212121/C2, which is why it + # survived). For a structure invariant under x -> Rx + t, and with the + # convention F(h) = Σ_j f_j exp(+2πi h·x_j): + # + # F(h) = Σ_j f_j exp(2πi h·(R x_j + t)) + # = exp(2πi h·t) · F(hR) + # => F(hR) = F(h) · exp(-2πi h·t) + # => arg F(h') - arg F(h) = -2π h·t for h' = hR = R^T h + # + # Verified against direct structure-factor summation and + # ``gemmi.Op.phase_shift``. The residual error of the wrong sign is + # 4π h·t mod 2π: zero for 2₁ screws and centring, π for 4₁/4₃, and + # 2π/3 for 3₁/6₁ -- see ``tests/unit/symmetry/test_phase_convention.py``. + phase_shift = -2.0 * np.pi * torch.matmul(hkl_float, translations[i]) all_hkl.append(hkl_transformed) all_phases.append(phase_shift) @@ -1267,7 +1281,8 @@ def get_canonical_hkl(hkl_single): t = translations[equiv_idx] hkl_trans = torch.round(torch.matmul(hkl_single, R.T)).to(torch.int32) - phase_shift = 2.0 * np.pi * torch.matmul(hkl_single, t) + # -2π h·t, same convention as expand_hkl (see the derivation there). + phase_shift = -2.0 * np.pi * torch.matmul(hkl_single, t) if tuple(hkl_trans.cpu().numpy()) == canonical: asu_reflections[canonical].append( @@ -1501,10 +1516,28 @@ def canonicalize_hkl( ) # --- Compute phase shifts vectorially --- - # phase_shift[i] = 2*pi * hkl_orig[i] . translations[op_idx[i]] + # The consumer applies these as + # phi_canonical = where(friedel_flags, -phi_input, phi_input) + phase_shifts + # so the required sign *depends on whether the row was Friedel-flipped*. + # With h' = hR (see the derivation in ``expand_hkl``), arg F(h') = arg F(h) - 2π h·t: + # + # non-Friedel, h_c = h R: + # phi_c = phi_h - 2π h·t => shift = -2π h·t + # Friedel, h_c = -(h R): + # phi_c = -phi_(hR) = -phi_h + 2π h·t + # and the consumer already supplies -phi_h => shift = +2π h·t + # + # A single uniform sign is therefore wrong for one of the two halves. This was + # +2π for both until 2026-07, which happened to be correct for the Friedel half + # and wrong for the rest; it is undetectable in P21/P212121/C2 because every + # shift there is 0 or π. See tests/unit/symmetry/test_phase_convention.py. t_selected = translations_np[op_idx] # (N, 3) + friedel_sign = np.where(friedel_np, 1.0, -1.0).astype(np.float32) phase_shifts_np = ( - 2.0 * np.pi * np.sum(hkl_np.astype(np.float32) * t_selected, axis=1) + friedel_sign + * 2.0 + * np.pi + * np.sum(hkl_np.astype(np.float32) * t_selected, axis=1) ).astype(np.float32) # --- Convert to tensors and sort --- diff --git a/torchref/symmetry/spacegroup.py b/torchref/symmetry/spacegroup.py index 02f4ecda..dda5067f 100644 --- a/torchref/symmetry/spacegroup.py +++ b/torchref/symmetry/spacegroup.py @@ -36,7 +36,7 @@ import torch import torch.nn as nn -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_float_dtype, normalize_device from torchref.utils.debug_utils import DebugMixin from torchref.utils.device_mixin import DeviceMovementMixin @@ -212,8 +212,7 @@ def get_operations_as_tensors( """ if dtype is None: dtype = get_float_dtype() - if device is None: - device = get_default_device() + device = normalize_device(device) sg = _normalize_spacegroup(spacegroup) # Extract rotation matrices and translations from gemmi operations @@ -692,8 +691,7 @@ def __init__( super(SpaceGroup, self).__init__() if dtype is None: dtype = get_float_dtype() - if device is None: - device = get_default_device() + device = normalize_device(device) self._device = device self._dtype = dtype @@ -864,9 +862,15 @@ def apply_to_hkl(self, hkl: torch.Tensor) -> torch.Tensor: """ Apply symmetry operations to Miller indices (rotation only, no translation). - For reciprocal space, only the rotational part of symmetry operations - applies to Miller indices: h' = R·h. The translation vector affects the - phase of structure factors, not the indices themselves. + For reciprocal space, the symmetry-equivalent Miller index is obtained + by the *transpose* of the (real-space, fractional) rotation matrix: + ``h' = h·R = Rᵀ·h``. This is NOT the same as the real-space transform + ``x' = R·x`` unless ``R`` is symmetric: for space groups whose fractional + rotation matrices are non-symmetric (trigonal, hexagonal, and permutation- + type cubic operations) applying ``R·h`` instead of ``Rᵀ·h`` yields the + wrong set of equivalents, corrupting centric-flag and epsilon-multiplicity + determination. The translation vector affects the phase of structure + factors, not the indices themselves, so it is not applied here. Parameters ---------- @@ -881,9 +885,12 @@ def apply_to_hkl(self, hkl: torch.Tensor) -> torch.Tensor: See Also -------- - apply : For real space coordinates (rotation + translation). + apply : For real space coordinates (rotation + translation), which uses + ``R·x`` and is the transpose of the reciprocal-space convention here. """ - return self.apply(hkl, apply_translation=False) + coords = hkl.to(self.matrices.device).to(self.matrices.dtype) + # result[n, i, o] = sum_j matrices[o, j, i] * coords[n, j] = (Rᵀ·h)_i + return torch.einsum("oji,nj->nio", self.matrices, coords) def expand_coords_to_P1(self, xyz_fractional: torch.Tensor) -> torch.Tensor: """ diff --git a/torchref/utils/__init__.py b/torchref/utils/__init__.py index 03db1064..10e0bd92 100644 --- a/torchref/utils/__init__.py +++ b/torchref/utils/__init__.py @@ -11,7 +11,7 @@ - Loss-finiteness validation (validate_loss) - Autograd introspection (collect_loss_leaves) - JSON serialization helpers (convert_to_serializable) -- Triton/eager backend dispatch (Engine) +- Backend dispatch (force_portable / use_portable) The names re-exported here are the package's public surface (see ``__all__``). Some submodule-only helpers (e.g. @@ -71,14 +71,12 @@ class MyRefinement(DebugMixin): # Serialization from .serialization import convert_to_serializable -# Triton/eager backend dispatch -from .triton_dispatch import ( - Engine, - get_engine, - set_engine, - should_use_triton, +# Backend dispatch +from .backends import ( + force_portable, + set_force_portable, triton_available, - use_engine, + use_portable, ) # Core utilities @@ -125,11 +123,9 @@ class MyRefinement(DebugMixin): "reset_diagnostic_budget", # Autograd introspection "collect_loss_leaves", - # Triton/eager backend dispatch - "Engine", - "get_engine", - "set_engine", - "use_engine", + # Backend dispatch + "force_portable", + "set_force_portable", + "use_portable", "triton_available", - "should_use_triton", ] diff --git a/torchref/utils/backends.py b/torchref/utils/backends.py new file mode 100644 index 00000000..559ab682 --- /dev/null +++ b/torchref/utils/backends.py @@ -0,0 +1,447 @@ +"""Declarative backend selection: the dispatch criteria, in one place, as data. + +Every dispatch site in TorchRef answers the same question -- given these tensors, which +kernel runs? This module holds the machinery for answering it from a table, so each criterion +(device, dtype, availability, failure policy) is written down exactly once, next to the +kernels it selects. +The tables themselves live with their kernels: see +``torchref/base/electron_density/_backends.py`` and +``torchref/base/direct_summation/_backends.py``. + +There is no chain +---------------- +The accelerator gates are pairwise device-disjoint -- CUDA, CPU and MPS -- so for any +given input **at most one** non-base backend can match. :func:`select` therefore returns a +single :class:`Backend`, not an ordered candidate list, and the only fallback is the +table's base case. An earlier design walked a chain; that models a control structure the +problem does not have. + +No configuration +---------------- +Selection takes no mode, tier or engine: for a given device and dtype there is exactly one +fastest usable kernel, so the answer is *derived*. The single override is ``force_portable``, +which pins the reference implementation -- see :func:`use_portable`. It exists for the one +failure automatic fallback cannot detect: an accelerator that runs and returns wrong numbers +rather than raising. + +This replaced a four-member selector enum whose two *forcing* members could never reach a +kernel the default would not already have picked -- they could only refuse when one was +unusable, which is information, not capability. + +Why kernels are stored as ``(module, attr)`` rather than as functions +-------------------------------------------------------------------- +Resolving ``getattr(import_module(path), attr)`` on every call costs a ``sys.modules`` hit +and a ``getattr``, and buys four things a captured function object would break: + +* ``torchref.utils`` is imported very early and must not reach into + ``torchref.base.electron_density`` at module scope; +* the Metal MSL source and ``import triton`` stay off hosts that never dispatch to them; +* availability stays readable *late*, so a test that flips + ``mps.compile._lib_failed`` is still seen; +* the kernel stays monkeypatchable at its defining module, which is what keeps the + provenance tests in ``tests/unit/structure_factor/test_dispatch.py`` meaningful. A table + holding function objects captured at import time would make those tests pass while + measuring nothing. +""" + +from __future__ import annotations + +import contextlib +import warnings +from dataclasses import dataclass, field +from importlib import import_module +from typing import Callable, Iterator, Optional, Sequence, Tuple + +import torch + + +__all__ = [ + "Backend", + "BackendTable", + "TorchRefDegradationWarning", + "force_portable", + "run_or_degrade", + "select", + "set_force_portable", + "use_portable", + "will_use", +] + + +# --------------------------------------------------------------------------- +# Forcing the portable reference kernel +# --------------------------------------------------------------------------- +# Named after the table row it selects (``portable``), so the flag and the backend cannot +# drift apart in the reader's head. +# +# One boolean is the whole override surface. Selection already picks the fastest backend that +# can run, and forcing an *accelerator* could never reach a kernel the default would not have +# picked -- it could only refuse when the kernel was unusable. What is genuinely useful is the +# other direction: pin the portable reference when you suspect an accelerator is producing +# wrong numbers, which no automatic fallback can detect. +_FORCE_PORTABLE: bool = False + + +def force_portable() -> bool: + """Whether dispatch is currently pinned to the portable reference kernel.""" + return _FORCE_PORTABLE + + +def set_force_portable(value: bool) -> None: + """Pin (or unpin) dispatch to the portable reference kernel, process-wide.""" + global _FORCE_PORTABLE + _FORCE_PORTABLE = bool(value) + + +@contextlib.contextmanager +def use_portable() -> Iterator[None]: + """Pin the portable reference kernel for the duration of the block. + + Restores the previous setting on exit, including on exception. + """ + previous = force_portable() + set_force_portable(True) + try: + yield + finally: + set_force_portable(previous) + + +class TorchRefDegradationWarning(UserWarning): + """An accelerator kernel failed at runtime and dispatch fell back to the reference. + + Its own category so tests can promote it to an error while production keeps running: a + fallback is a performance problem for a user and a correctness signal for CI. Before this + existed the fallback was entirely silent, and the only way to discover it was to force the + accelerator and watch it raise -- which meant you had to already suspect the problem. + """ + + +# --------------------------------------------------------------------------- +# Availability +# --------------------------------------------------------------------------- +_has_triton: Optional[bool] = None + + +def triton_available() -> bool: + """Whether ``triton`` is importable (cheap, cached for the process). + + Lives here beside the per-backend probes that consume it. ``except Exception``, not + ``except ImportError``: a Triton that is installed but skewed against the driver or LLVM + raises something else, and that must read as "unavailable" rather than propagating. + """ + global _has_triton + if _has_triton is None: + try: + import triton # noqa: F401 + + _has_triton = True + except Exception: + _has_triton = False + return _has_triton + + +def _mps_present() -> bool: + return hasattr(torch.backends, "mps") and torch.backends.mps.is_available() + + +#: Host conditions a backend can declare via ``Backend.expect_available``. Parametrized +#: rather than hand-written per kernel: "this must work here" is the same question for the +#: fused C++ splat, the Metal shader and the Triton kernels, and only the condition differs. +_EXPECTATIONS = { + "never": lambda: False, + "always": lambda: True, + "cuda": torch.cuda.is_available, + "mps": _mps_present, +} + + +@dataclass(frozen=True) +class Backend: + """One row of a dispatch table: a kernel plus every criterion for choosing it. + + Parameters + ---------- + name : str + Stable identifier, used in error messages and by tests to name a path. + kernel : (str, str, str), optional + ``(module_path, isotropic_attr, anisotropic_attr)``, resolved per call. For + single-variant backends pass the same attribute twice. + + ``None`` marks a **gate-only** backend: one whose selection criteria are worth + declaring in a table even though the call site names its own kernel. The geometry + targets are the case -- twelve modules each with a single ``if use_triton(x): from + .triton.X import f`` call site, where the three-line local import is more legible + than a registry hop, but "which devices and dtypes admit Triton here" still needs to + be stated once. :func:`run_or_degrade` refuses such a backend; only :func:`select` and + :func:`will_use` accept it. + device : str, optional + Required device type (``"cuda"``/``"mps"``/``"cpu"``). ``None`` means any, which + is what marks the base case. + dtypes : tuple[torch.dtype, ...], optional + Permitted **floating-point** dtypes. ``None`` means any. + + Integer tensors are exempt, deliberately. Miller indices arrive as ``int32`` from + the MTZ reader and every kernel casts them itself; that cast is exact for + \\|h\\| < 2**24, so an identity test against ``float32`` would reject the + production dtype and disable the kernel it was meant to protect. What the rule + catches is a *float* in the wrong precision -- a float64 ``hkl`` whose silent + downcast would truncate the phase. + require_uniform_dtype : bool + Whether every probed floating-point tensor must share **one** dtype, as opposed to + each independently being in ``dtypes``. + + Not the same condition, and the difference is memory safety rather than taste. The + fused CPU kernel selects one ``scalar_t`` from the output map via + ``AT_DISPATCH_FLOATING_TYPES`` and then reads every other tensor through a raw + pointer of that type, so a float64 map beside float32 atoms reinterprets the + coordinate buffer as doubles -- a 2x out-of-bounds read. ``dtypes=(f32, f64)`` + alone *admits* that call. + probes : tuple[int, ...], optional + Which argument positions carry the device/dtype contract. ``None`` probes them all. + + Per-table because the requirement means different things in different tables. For + the density splats float32 is a *capability* -- the C++ and the MSL shader cannot + consume anything else. For direct summation it is *policy*: the Triton kernel casts + everything itself, so the gate polices only the leaves whose precision the caller + chose, and must not probe the float32 stored scattering-factor table (which would + decline the kernel unconditionally). + probe : (str, str), optional + ``(module_path, attr)`` of a zero-argument callable returning ``None`` when the + backend is usable, or a human-readable reason when it is not. ``None`` means + always available. Resolved late, so a cache the callable consults can still be + invalidated. + + One function per backend, deliberately: a separate boolean ``*_available()`` would + be the same test written twice. Where such a bool already exists as public API it + derives from this probe rather than repeating it. + expect_available : {"never", "always", "cuda", "mps"} + The host condition under which this backend *must* work. Distinct from ``probe``, + which reports what is true; this states what ought to be. + + The two together are what let a broken build fail rather than skip. Dispatch under + ``AUTO`` degrades quietly, which is right for users and dangerous for CI: if every + test skips when a kernel is missing, a build that stopped working produces an + all-green run while production has silently fallen back. Declaring the expectation + here means one parametrized test covers every backend, instead of each kernel + needing its own bespoke compile check. + + Not consulted by :func:`select` -- availability is a fact at dispatch time, never an + expectation. + on_failure : {"raise", "degrade"} + What a *runtime* exception from the kernel means: ``"degrade"`` falls back to the + base case (with a warning), ``"raise"`` propagates. + + ``"degrade"`` carries a precondition: the kernel must not have mutated its inputs + before failing, or the fallback would double-count. Both accelerator splats satisfy + it by cloning the density map before accumulating. + second_order : bool + Whether the kernel composes under ``create_graph=True``. Not consulted by + :func:`select`; it is here so the test matrix can be derived from this table rather + than maintaining its own copy. + """ + + name: str + kernel: Optional[Tuple[str, str, str]] + device: Optional[str] = None + dtypes: Optional[Tuple[torch.dtype, ...]] = None + require_uniform_dtype: bool = False + probes: Optional[Tuple[int, ...]] = None + probe: Optional[Tuple[str, str]] = None + expect_available: str = "never" + on_failure: str = "raise" + second_order: bool = True + + def __post_init__(self): + if self.on_failure not in ("raise", "degrade"): + raise ValueError( + f"{self.name}: on_failure must be 'raise' or 'degrade', " + f"got {self.on_failure!r}" + ) + if self.expect_available not in _EXPECTATIONS: + raise ValueError( + f"{self.name}: expect_available must be one of " + f"{sorted(_EXPECTATIONS)}, got {self.expect_available!r}" + ) + + def expected_here(self) -> bool: + """Whether this host is one where this backend is required to work.""" + return _EXPECTATIONS[self.expect_available]() + + # -- criteria --------------------------------------------------------- + def requirement(self) -> str: + """The device/dtype contract as an error fragment, e.g. ``requires MPS float32``.""" + parts = [] + if self.device is not None: + parts.append(self.device.upper()) + if self.dtypes is not None: + parts.append("/".join(str(d).replace("torch.", "") for d in self.dtypes)) + if not parts: + return "accepts any inputs" + return "requires " + " ".join(parts) + " inputs" + + def mismatch(self, tensors: Sequence[Optional[torch.Tensor]]) -> Optional[str]: + """Why these tensors fail this backend's device/dtype contract, or ``None``.""" + probed = _probed(self, tensors) + if self.device is not None and any( + t.device.type != self.device for t in probed + ): + return self.requirement() + if self.dtypes is not None: + floats = [t for t in probed if t.is_floating_point()] + if any(t.dtype not in self.dtypes for t in floats): + return self.requirement() + if self.require_uniform_dtype and len({t.dtype for t in floats}) > 1: + return ( + f"{self.requirement()} sharing a single dtype (got " + + ", ".join( + sorted(str(d).replace("torch.", "") for d in {t.dtype for t in floats}) + ) + + ")" + ) + return None + + def unavailable(self) -> Optional[str]: + """Why this backend cannot run on this host, or ``None``.""" + if self.probe is None: + return None + module, attr = self.probe + try: + fn = getattr(import_module(module), attr) + except Exception as exc: # noqa: BLE001 - a stripped install must not crash a gate + return f"its availability probe could not be imported ({type(exc).__name__}: {exc})" + return fn() + + def resolve(self, aniso: bool) -> Callable: + """The kernel function, looked up now rather than captured at import.""" + if self.kernel is None: + raise TypeError( + f"{self.name} is gate-only (kernel=None): its call site names its own " + "kernel, so there is nothing here to run." + ) + module, iso_attr, aniso_attr = self.kernel + return getattr(import_module(module), aniso_attr if aniso else iso_attr) + + +@dataclass(frozen=True) +class BackendTable: + """An ordered set of backends plus the invariant that makes it a total policy. + + Checked at import: exactly one base case, i.e. one backend with no device and no dtype + restriction. That is what makes :func:`select` unable to fail -- with no unrestricted row + there would be inputs no backend matched, and selection would need an error path. + """ + + name: str + backends: Tuple[Backend, ...] + base: Backend = field(init=False) + + def __post_init__(self): + bases = [b for b in self.backends if b.device is None and b.dtypes is None] + if len(bases) != 1: + raise ValueError( + f"{self.name}: expected exactly one unrestricted base backend, " + f"found {[b.name for b in bases]}" + ) + object.__setattr__(self, "base", bases[0]) + + def by_name(self, name: str) -> Backend: + for b in self.backends: + if b.name == name: + return b + raise KeyError(f"{self.name}: no backend named {name!r}") + + +def _probed(backend: Backend, tensors: Sequence[Optional[torch.Tensor]]): + """The tensors this backend's contract applies to: selected, then ``None``-filtered.""" + if backend.probes is None: + chosen = tensors + else: + chosen = [tensors[i] for i in backend.probes if i < len(tensors)] + return [t for t in chosen if t is not None] + + +def select( + table: BackendTable, + tensors: Sequence[Optional[torch.Tensor]], + force_portable: Optional[bool] = None, +) -> Backend: + """The one backend that runs. Cannot fail. + + Totality is structural, not defensive: the base case declares no device and no dtype and + carries no probe, so both its criteria return ``None`` unconditionally and the loop always + terminates there. There is nothing to raise about -- selection is asked "what is fastest + here", and the answer is never "nothing". + + Two-phase by contract: device and dtype are checked for every candidate *before* any + availability probe is called. That ordering is load-bearing -- it keeps an MPS host from + compiling the CPU C++ extension it will never use, and a CUDA host from importing Triton + to answer a question about a CPU tensor. + + Parameters + ---------- + table : BackendTable + The policy to apply. + tensors : sequence of torch.Tensor or None + Positional arguments to probe. ``None`` entries are ignored, so a site may pass + optional inputs straight through. + force_portable : bool, optional + Pin the portable reference kernel. ``None`` defers to the process-wide setting. + """ + pin = force_portable if force_portable is not None else _FORCE_PORTABLE + if pin: + return table.base + + for backend in table.backends: + if backend.mismatch(tensors) is None and backend.unavailable() is None: + return backend + return table.base + + +def will_use( + table: BackendTable, + name: str, + tensors: Sequence[Optional[torch.Tensor]], + force_portable: Optional[bool] = None, +) -> bool: + """Whether ``name`` is the backend that would actually run for these inputs. + + A thin wrapper over :func:`select`: with nothing to force there is no third answer, so + the question reduces to "is the selected backend this one". + """ + return select(table, tensors, force_portable=force_portable).name == name + + +def run_or_degrade( + table: BackendTable, + backend: Backend, + aniso: bool, + *args, + **kwargs, +): + """Run ``backend``'s kernel, falling back to the table's base case only if allowed. + + ``on_failure`` alone decides. A backend marked ``"degrade"`` falls back, which keeps a + misbehaving accelerator a performance problem rather than an outage; one marked + ``"raise"`` propagates, because a kernel that built fine and then threw is a bug, and + substituting the reference would convert wrong results into a large silent slowdown whose + numbers still look plausible. + + A fallback always **warns**. It is a performance problem for a user and a correctness + signal for CI, and it used to be entirely silent -- discoverable only by already + suspecting it and forcing the accelerator to watch it raise. + """ + fn = backend.resolve(aniso) + if backend.on_failure == "raise" or backend is table.base: + return fn(*args, **kwargs) + try: + return fn(*args, **kwargs) + except Exception as exc: # noqa: BLE001 - the fallback is the point + warnings.warn( + f"{table.name}: the {backend.name} kernel failed and dispatch fell back to " + f"{table.base.name} ({type(exc).__name__}: {exc}). Results are correct but " + "slower; this is worth investigating.", + TorchRefDegradationWarning, + stacklevel=2, + ) + return table.base.resolve(aniso)(*args, **kwargs) diff --git a/torchref/utils/device_mixin.py b/torchref/utils/device_mixin.py index f3533305..72f8ecf4 100644 --- a/torchref/utils/device_mixin.py +++ b/torchref/utils/device_mixin.py @@ -35,11 +35,33 @@ class MyDataclass(DeviceMixin): from __future__ import annotations +import inspect import threading import torch from torch import nn +from torchref.config import canonical_device + +# ``torch._C._nn._parse_to`` is the same parser ``nn.Module.to`` uses, so it +# accepts every overload PyTorch does -- including ``.to(other_tensor)`` and +# ``.to(0)``, which the hand-rolled ``_parse_to_args`` fallback below cannot +# express. It is a private API, so probe once at import rather than catching +# per call: catching ``TypeError``/``RuntimeError`` around each invocation +# would turn a genuinely invalid user argument into a silent no-op. +try: + _PARSE_TO = torch._C._nn._parse_to +except AttributeError: # pragma: no cover - defensive against API drift + _PARSE_TO = None + +# Whether this torch exposes ``nn.Module._apply(fn, recurse=...)``. Detected +# once by signature rather than by catching ``TypeError`` at the call site: a +# blanket catch there would also swallow a genuine ``TypeError`` raised *inside* +# ``_apply`` after some tensors had already moved, and retry the move. +_MODULE_APPLY_TAKES_RECURSE = ( + "recurse" in inspect.signature(nn.Module._apply).parameters +) + # --------------------------------------------------------------------------- # Thread-local traversal state (cycle detection across one top-level .to()) # --------------------------------------------------------------------------- @@ -71,11 +93,50 @@ class MyDataclass(DeviceMixin): ) +class _ToRequest: + """The ``(device, dtype)`` a top-level ``.to()`` asked for. + + :meth:`DeviceMixin._apply` receives only an opaque tensor-to-tensor ``fn``, + so an object that owns no tensors cannot work out where it was just asked + to go. Recording the parsed request on the traversal thread-local lets + :func:`_refresh_device_trackers` answer that question at *every* node -- + including tensor-free children, which are reached through + ``child._apply(fn)`` and therefore never pass through their own ``to()``. + + ``probe`` caches a :func:`_probe_target` result for the traversal so a + graph with many tensor-free nodes probes at most once. + """ + + __slots__ = ("device", "dtype", "probe") + + def __init__(self, device=None, dtype=None): + self.device = canonical_device(device) + self.dtype = dtype + self.probe = None + + def _current_visited(): """Return the active visited set, or ``None`` if not in a traversal.""" return getattr(_traversal_state, "visited", None) +def _current_request(): + """Return the active :class:`_ToRequest`, or ``None``.""" + return getattr(_traversal_state, "request", None) + + +def _push_request(request): + """Install ``request`` for the current traversal, returning the previous.""" + prev = getattr(_traversal_state, "request", None) + _traversal_state.request = request + return prev + + +def _pop_request(prev): + """Restore the request saved by :func:`_push_request`.""" + _traversal_state.request = prev + + def _enter_traversal(): """Begin a top-level traversal. Returns a token for :func:`_exit_traversal`. @@ -85,6 +146,9 @@ def _enter_traversal(): prev = getattr(_traversal_state, "visited", None) if prev is None: _traversal_state.visited = set() + # Never inherit a request leaked by an earlier traversal that unwound + # abnormally: a stale target would silently misdirect this one. + _traversal_state.request = None return "owner" return None @@ -93,6 +157,7 @@ def _exit_traversal(token): """End a top-level traversal.""" if token == "owner": _traversal_state.visited = None + _traversal_state.request = None def _parse_to_args(args, kwargs): @@ -116,6 +181,15 @@ def _parse_to_args(args, kwargs): return device, dtype +def _accepts_single_arg(func) -> bool: + """Whether ``func`` can be called with exactly one positional argument.""" + try: + inspect.signature(func).bind(None) + except (TypeError, ValueError): + return False + return True + + def _apply_to_obj(val, fn, visited): """Apply ``fn`` to any tensors inside ``val``, recursing as needed. @@ -147,12 +221,13 @@ def _apply_to_obj(val, fn, visited): if callable(apply_method) and not isinstance(val, type): if id(val) in visited: return val - try: + # Check the signature up front instead of catching ``TypeError`` from + # the call: a catch there cannot distinguish "wrong signature" from a + # real ``TypeError`` raised mid-traversal, and would silently leave the + # object partially moved. + if _accepts_single_arg(apply_method): apply_method(fn) return val - except TypeError: - # Object has an _apply with a different signature; skip silently. - pass if isinstance(val, list): new_list = [_apply_to_obj(v, fn, visited) for v in val] @@ -180,68 +255,259 @@ def _apply_to_obj(val, fn, visited): def _invalidate_caches(obj): - """Call ``reset_forward_cache`` / ``reset_cache`` if present.""" - reset_fwd = getattr(obj, "reset_forward_cache", None) - if callable(reset_fwd): - try: - reset_fwd() - except Exception: - pass - reset_full = getattr(obj, "reset_cache", None) - if callable(reset_full): + """Call ``reset_forward_cache`` / ``reset_cache`` if present. + + Failures propagate rather than being swallowed. A cache that fails to clear + keeps tensors from the *previous* device, which produces either a + cross-device error much later or -- worse -- silently stale numbers; that + is precisely the corruption this mixin exists to prevent, so it must not be + hidden behind a bare ``except``. + + Raises + ------ + RuntimeError + Chained from the hook's own exception, naming the object and hook. + Note that movement is **not atomic**: by the time a cache hook runs, + some of the object's tensors have already been transformed, so the + object may be left partially moved. + """ + for hook_name in ("reset_forward_cache", "reset_cache"): + hook = getattr(obj, hook_name, None) + if not callable(hook): + continue try: - reset_full() - except Exception: - pass - - -def _representative_tensor(obj): - """Return one tensor that reflects the current device/dtype of *obj*. - - Prefers buffers (their dtype tracks ``dtype_float`` for crystallographic - code) and falls back to parameters, then to a plain ``torch.Tensor`` - attribute found in ``__dict__``. + hook() + except Exception as exc: + raise RuntimeError( + f"{type(obj).__name__}.{hook_name}() failed during device/dtype " + f"movement: {type(exc).__name__}: {exc}. The object may be " + "partially moved and its caches may still hold tensors on the " + "previous device." + ) from exc + + +def _owned_tensors(obj): + """Yield the tensors *obj* itself owns (never a child's). + + ``recurse=False`` is deliberate: the trackers describe where this object + allocates, so inheriting a device from a submodule would report a + half-moved graph as consistent. """ if isinstance(obj, nn.Module): - for buf in obj.buffers(): - return buf - for param in obj.parameters(): - return param + # Read the registration dicts rather than ``buffers()``/``parameters()``: + # several classes here override those accessors with a no-argument + # signature (``SolventModel.parameters``, ``MixedTensor.parameters``), + # and the dicts are in any case the literal "owned, not inherited" set. + for buf in obj._buffers.values(): + if buf is not None: + yield buf + for param in obj._parameters.values(): + if param is not None: + yield param for val in obj.__dict__.values(): if isinstance(val, torch.Tensor): - return val - if hasattr(val, "_data") and isinstance(getattr(val, "_data"), torch.Tensor): - return val._data + yield val + else: + data = getattr(val, "_data", None) + if isinstance(data, torch.Tensor): + yield data + # Container subclasses (``TensorMasks`` is a ``dict``) keep their tensors in + # the container's own storage, not in ``__dict__``. + if isinstance(obj, dict): + for val in obj.values(): + if isinstance(val, torch.Tensor): + yield val + elif isinstance(obj, (list, tuple)): + for val in obj: + if isinstance(val, torch.Tensor): + yield val + + +def _representative_tensor(obj): + """Return one tensor reflecting *obj*'s own device, or ``None``.""" + for tensor in _owned_tensors(obj): + return tensor return None -def _refresh_device_trackers(obj): +def _observed_state(obj): + """Return ``(device, floating_dtype)`` observed from *obj*'s own tensors. + + The two axes are resolved **independently**: the first owned tensor fixes + the device, but only a floating/complex tensor may fix ``dtype_float``. + Deciding both from a single representative tensor lets an integer buffer + that happens to be registered first (``hkl``, ``aniso_flag``) veto the + dtype answer. + """ + device = None + dtype = None + for tensor in _owned_tensors(obj): + if device is None: + device = tensor.device + if dtype is None and (tensor.is_floating_point() or tensor.is_complex()): + dtype = tensor.dtype + if device is not None and dtype is not None: + break + return device, dtype + + +def _probe_target(fn): + """Discover what ``fn`` does to a tensor's device and floating dtype. + + Needed only on the paths that bypass :meth:`DeviceMixin.to` and therefore + record no request: ``.float()`` / ``.double()`` / ``.half()``, and an + ``_apply`` driven by a plain (non-mixin) ``nn.Module`` parent. + + ``fn`` is probed on scratch tensors whose devices are **known to be + valid**, never on the object's own tracker -- a tensor-free object can be + carrying a stale tracker naming a backend this host does not have (say + ``cuda:0`` on a CPU-only machine), and allocating there would fail exactly + when the answer matters most. + + **Both** axes need two reference points, and the scratches must differ on + both. A transforming ``fn`` funnels its two inputs to a single value on the + axis it touches, while a preserving ``fn`` hands each input's own value + back: + + ================== =================== =================== + ``fn`` devices agree? dtypes agree? + ================== =================== =================== + ``.to(device=D)`` yes -> device is D no -> dtype is None + ``.half()`` no -> None yes -> dtype is f16 + ``.to(D, f32)`` yes -> D yes -> f32 + ================== =================== =================== + + Probing one axis with a single reference is what makes a device-only move + report the scratch's own dtype and clobber ``dtype_float``. + + Scratch devices are always ones **known to be valid**, never the object's + own tracker -- a tensor-free object can carry a stale tracker naming a + backend this host does not have (``cuda:0`` on a CPU-only machine), and + allocating there would fail exactly when the answer matters most. + + Returns + ------- + tuple + ``(device, dtype)``; either is ``None`` when ``fn`` leaves that axis + untouched, or when ``fn`` could not be probed at all. + """ + accel = None + if torch.cuda.is_available(): + accel = torch.device("cuda", torch.cuda.current_device()) + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + accel = torch.device("mps", 0) + + def run(device, dtype): + """Apply ``fn`` to one scratch tensor; ``None`` if that is not possible.""" + try: + scratch = torch.empty(0, device=device, dtype=dtype) + except (RuntimeError, TypeError): + return None # backend advertised but unusable for this dtype + # ``fn`` is arbitrary caller-supplied code, so a broad catch is correct + # here -- unlike the cache hooks above, a non-conforming ``fn`` is an + # expected outcome, not a corruption signal. The caller decides what an + # unprobeable ``fn`` means. + try: + out = fn(scratch) + except Exception: + return None + return out if isinstance(out, torch.Tensor) else None + + # The two axes get their own contrast pair, and each pair varies only along + # the axis it measures. Sharing one pair for both would couple them: an + # accelerator scratch cannot be cast to float64 on MPS, so probing dtype on + # the device pair makes ``.double()`` unprobeable on Apple silicon. + base = run(torch.device("cpu"), torch.float32) + if base is None: + return None, None + + device = base.device + if accel is not None: + # Contrast pair for the device axis: same dtype, different device. + other = run(accel, torch.float32) + if other is None: + device = None + elif other.device != base.device: + device = None # each kept its own device -> device-preserving + + dtype = None + if base.is_floating_point() or base.is_complex(): + # Contrast pair for the dtype axis: same device, different dtype. + # float16 rather than float64 so this stays cheap and universally + # supported; the CPU pin means ``.double()`` remains probeable. + other = run(torch.device("cpu"), torch.float16) + if other is not None and other.dtype == base.dtype: + dtype = base.dtype + + return device, dtype + + +_DEVICE_TRACKERS = ("device", "_device") +_DTYPE_TRACKERS = ("dtype_float", "_dtype") + + +def _refresh_device_trackers(obj, fn=None): """Refresh ``device`` / ``_device`` / ``dtype_float`` / ``_dtype`` trackers. - Only updates attributes whose current value is already a ``torch.device`` - / ``torch.dtype`` (or ``None`` / ``str`` / ``int``) — never overwrites - unrelated attributes that happen to share the same name. The tracker - controls where new tensors are subsequently allocated, so it must follow - ``.to()`` moves. + The tracker decides where the object allocates *next* -- over 200 call + sites in the package pass ``device=self.device`` -- so leaving it stale on + a tensor-free object silently misplaces every tensor that object later + creates. + + Resolution order, applied per axis: + + 1. a tensor the object itself owns (authoritative, already moved), + 2. the ``.to()`` request recorded on the traversal thread-local, which is + the only source available to an object holding no tensors, + 3. a probe of ``fn`` (see :func:`_probe_target`), for the ``.float()`` and + plain-parent paths that record no request. + + Only attributes already present in ``obj.__dict__`` and already holding a + device/dtype-shaped value are written, so an unrelated attribute that + happens to be named ``device`` is never clobbered. Written values are + canonical, so ``obj.device == some_tensor.device`` holds. """ - rep = _representative_tensor(obj) - if rep is None: + has_device_tracker = any(a in obj.__dict__ for a in _DEVICE_TRACKERS) + has_dtype_tracker = any(a in obj.__dict__ for a in _DTYPE_TRACKERS) + if not (has_device_tracker or has_dtype_tracker): return - for attr_name in ("device", "_device"): - if attr_name not in obj.__dict__: - continue - current = obj.__dict__[attr_name] - if current is None or isinstance(current, (torch.device, str, int)): - obj.__dict__[attr_name] = rep.device + device, dtype = _observed_state(obj) + + if device is None or dtype is None: + request = _current_request() + if request is not None: + if device is None: + device = request.device + if dtype is None: + dtype = request.dtype + if (device is None or dtype is None) and fn is not None: + if request.probe is None: + request.probe = _probe_target(fn) + probed_device, probed_dtype = request.probe + device = device if device is not None else probed_device + dtype = dtype if dtype is not None else probed_dtype + elif fn is not None: + probed_device, probed_dtype = _probe_target(fn) + device = device if device is not None else probed_device + dtype = dtype if dtype is not None else probed_dtype + + if device is not None: + device = canonical_device(device) + for attr_name in _DEVICE_TRACKERS: + if attr_name not in obj.__dict__: + continue + current = obj.__dict__[attr_name] + if current is None or isinstance(current, (torch.device, str, int)): + obj.__dict__[attr_name] = device - if rep.is_floating_point() or rep.is_complex(): - for attr_name in ("dtype_float", "_dtype"): + if dtype is not None: + for attr_name in _DTYPE_TRACKERS: if attr_name not in obj.__dict__: continue current = obj.__dict__[attr_name] if current is None or isinstance(current, torch.dtype): - obj.__dict__[attr_name] = rep.dtype + obj.__dict__[attr_name] = dtype def _safe_setattr(obj, name, value): @@ -293,19 +559,39 @@ def to(self, *args, **kwargs): # type: ignore[override] pair is parsed and applied via :meth:`_apply`. A call that resolves to neither a device nor a dtype is a no-op that returns ``self``. """ + if _PARSE_TO is not None: + # Let PyTorch's own parser raise on invalid arguments. + device, dtype = _PARSE_TO(*args, **kwargs)[:2] + else: # pragma: no cover - only on a torch build without the private API + device, dtype = _parse_to_args(args, kwargs) + + # Build the request BEFORE claiming the traversal. ``_ToRequest`` + # canonicalises the device, which raises for a backend this host does + # not have (``.to('cuda')`` on a CUDA-less machine). Constructing it + # after ``_enter_traversal`` but outside the ``try`` leaked the + # thread-local visited set on that raise, and every later ``.to()`` on + # the thread then short-circuited as "already visited" -- moving + # nothing, silently. + request = _ToRequest(device, dtype) + token = _enter_traversal() + prev_request = _push_request(request) try: if isinstance(self, nn.Module): return super().to(*args, **kwargs) - device, dtype = _parse_to_args(args, kwargs) if device is None and dtype is None: return self + # Forward the caller's original arguments rather than the parsed + # pair, so overloads and options the parser folds away -- + # ``.to(other_tensor)``, ``non_blocking=``, ``memory_format=`` -- + # reach the tensors intact. def fn(t): - return t.to(device=device, dtype=dtype) + return t.to(*args, **kwargs) return self._apply(fn) finally: + _pop_request(prev_request) _exit_traversal(token) def cuda(self, device=None): # type: ignore[override] @@ -356,11 +642,15 @@ def _apply(self, fn, recurse=True): # type: ignore[override] return self visited.add(id(self)) + # Snapshot before anything moves, so step 5 can tell a real + # transformation from a ``.to()`` that targets the current state. + old_device, old_dtype = _observed_state(self) + # 1. Standard nn.Module traversal (params, buffers, child modules). if isinstance(self, nn.Module): - try: + if _MODULE_APPLY_TAKES_RECURSE: super()._apply(fn, recurse=recurse) - except TypeError: + else: # pragma: no cover - older torch without the kwarg super()._apply(fn) # Mark registered children as visited so that other plain @@ -379,15 +669,72 @@ def _apply(self, fn, recurse=True): # type: ignore[override] # 3. Refresh ``device`` / ``_device`` / ``dtype`` trackers so # subsequent tensor allocations target the new device. - _refresh_device_trackers(self) + _refresh_device_trackers(self, fn) # 4. Invalidate caches. _invalidate_caches(self) + + # 5. Notify the object iff something actually changed, so index + # rebuilds do not fire on a no-op ``.to(current_device)``. + new_device, new_dtype = _observed_state(self) + device_changed = ( + old_device is not None + and new_device is not None + and canonical_device(old_device) != canonical_device(new_device) + ) + dtype_changed = ( + old_dtype is not None + and new_dtype is not None + and old_dtype != new_dtype + ) + if device_changed or dtype_changed: + self._after_device_apply( + old_device, + new_device, + old_dtype, + new_dtype, + device_changed=device_changed, + dtype_changed=dtype_changed, + ) return self finally: if token is not None: _exit_traversal(token) + def _after_device_apply( + self, + old_device, + new_device, + old_dtype, + new_dtype, + *, + device_changed, + dtype_changed, + ): + """Hook called once per object after a *real* device/dtype change. + + No-op by default. Override to rebuild state derived from tensor + placement -- precomputed index tensors, device-specific caches -- that + the generic walk cannot repair on its own. + + Use this rather than ``reset_cache()``: ``reset_cache`` is a + functional-cache hook that :meth:`LossState.reset_caches` fires after + *every* optimizer step, so rebuilding indices there would put index + reconstruction and GPU syncs on the hot path. This hook only fires when + movement actually happened. + + Parameters + ---------- + old_device, new_device : torch.device or None + Device before and after. ``None`` when the object owned no tensors + on that side of the move. + old_dtype, new_dtype : torch.dtype or None + Floating/complex dtype before and after, resolved independently of + the device. + device_changed, dtype_changed : bool + Which axes actually changed. At least one is always True. + """ + # --------------------------------------------------------------------------- # Backwards-compatibility aliases diff --git a/torchref/utils/device_resolution.py b/torchref/utils/device_resolution.py index f9680023..8029439d 100644 --- a/torchref/utils/device_resolution.py +++ b/torchref/utils/device_resolution.py @@ -18,19 +18,10 @@ import torch +from torchref.config import canonical_device as _canonical from torchref.config import get_default_device -def _canonical(device: torch.device) -> torch.device: - """Return ``device`` with its default index filled in. - - ``torch.device('cuda') != torch.device('cuda:0')`` even though both - point to the same physical device. Materialising an empty tensor - on the device is the cheapest way to get the canonical form. - """ - return torch.empty(0, device=device).device - - def resolve_device( *modules: Any, device: Optional[Union[torch.device, str]] = None, @@ -71,6 +62,17 @@ def resolve_device( torch.device The resolved device. + Raises + ------ + TypeError + If a bare ``torch.Tensor`` (or ``nn.Parameter``) is passed. Such an + object satisfies the ``.device`` / ``.to()`` precondition + *syntactically* but violates it semantically, because + ``Tensor.to()`` returns a new tensor rather than moving in place -- + so the move would be dropped without a word. Use + :func:`torchref.config.normalize_device` to read a device off a + tensor; use this function only to reconcile owning objects. + Examples -------- Empty call returns the configured default:: @@ -89,10 +91,29 @@ def resolve_device( >>> resolve_device(cuda_model, cpu_data) # doctest: +SKIP device(type='cuda') """ + for m in modules: + if isinstance(m, torch.Tensor): + raise TypeError( + "resolve_device() moves its inputs in place and returns only " + f"the resolved device, but a bare {type(m).__name__} was " + "passed. torch.Tensor.to() is out-of-place, so the move would " + "be silently discarded and the tensor left where it was. Pass " + "the object that owns the tensor, or move it yourself: " + "t = t.to(normalize_device(...))." + ) + if device is not None: - resolved = torch.device(device) if not isinstance(device, torch.device) else device + resolved = _canonical(device) for m in modules: - if m is not None: + # Skip modules already on target. ``.to()`` invalidates caches and + # fires ``_after_device_apply``, so a no-op move is not free -- + # ``SfDS.forward`` calls ``resolve_device`` on every evaluation. + # + # This is only safe because ``DeviceMixin`` now keeps ``m.device`` + # truthful; before that, the unconditional ``.to()`` here was + # accidentally repairing objects whose tracker lied about where + # their tensors were. + if m is not None and _canonical(m.device) != resolved: m.to(resolved) return resolved diff --git a/torchref/utils/triton_dispatch.py b/torchref/utils/triton_dispatch.py deleted file mode 100644 index f6d2dc2b..00000000 --- a/torchref/utils/triton_dispatch.py +++ /dev/null @@ -1,153 +0,0 @@ -"""Shared Triton/eager backend dispatch for TorchRef. - -A single capability-based selector used by every dispatch site (direct -summation, geometry/X-ray targets, FFT electron density). For a given -``(device, dtype)`` there is exactly one best path, so the backend is -*derived* rather than configured: - -- CUDA + float32 + Triton available -> Triton kernel -- everything else (CPU, float64, MPS, no Triton) -> eager fallback - -The :class:`Engine` enum exists only to *force* a path for tests and -benchmarks. There is **no environment-variable dispatch** — selection is via -the engine object (process-wide default + per-call override) only. - -Examples --------- -:: - - from torchref.utils import Engine, use_engine, should_use_triton - - # force eager for an A/B block (e.g. in a test) - with use_engine(Engine.EAGER): - loss = target() - - # inside a dispatch site - if should_use_triton(xyz): # reads the global engine - return _triton_path(...) - return _eager_path(...) -""" - -from __future__ import annotations - -import contextlib -import enum -from typing import Iterator, Optional - -import torch - -__all__ = [ - "Engine", - "get_engine", - "set_engine", - "use_engine", - "triton_available", - "should_use_triton", -] - - -class Engine(enum.Enum): - """Backend selector shared by all Triton dispatch sites. - - ``AUTO`` derives the backend from device/dtype/availability. ``TRITON`` - and ``EAGER`` force a path; ``TRITON`` raises if CUDA+float32+availability - is not met (force = strict, never silently degrade). - """ - - AUTO = "auto" - TRITON = "triton" - EAGER = "eager" - - -# Process-wide default engine. No environment-variable read — change via -# set_engine()/use_engine() or a per-call engine= override. -_ENGINE: Engine = Engine.AUTO - - -def get_engine() -> Engine: - """Return the current process-wide engine.""" - return _ENGINE - - -def set_engine(engine: Engine) -> None: - """Set the process-wide engine.""" - global _ENGINE - _ENGINE = engine - - -@contextlib.contextmanager -def use_engine(engine: Engine) -> Iterator[None]: - """Temporarily force an engine for the duration of the ``with`` block. - - Restores the previous engine on exit (including on exception). - """ - prev = get_engine() - set_engine(engine) - try: - yield - finally: - set_engine(prev) - - -_has_triton: Optional[bool] = None - - -def triton_available() -> bool: - """Whether ``triton`` is importable (cheap, cached for the process).""" - global _has_triton - if _has_triton is None: - try: - import triton # noqa: F401 - - _has_triton = True - except Exception: - _has_triton = False - return _has_triton - - -def should_use_triton(*tensors: torch.Tensor, engine: Optional[Engine] = None) -> bool: - """Coarse triton-vs-eager gate shared by every dispatch site. - - Reads the process-wide engine unless an explicit ``engine`` is passed - (direct summation / ``SfDS`` use the per-call override). Probes the given - tensors: the Triton path requires every non-None tensor to be CUDA - float32. - - Parameters - ---------- - *tensors : torch.Tensor - Tensors whose device/dtype gate the Triton path. ``None`` entries are - ignored (a target may pass optional tensors). - engine : Engine, optional - Per-call override. Defaults to the process-wide engine. - - Returns - ------- - bool - True if the Triton kernel should be used. - - Raises - ------ - RuntimeError - If ``engine`` is ``Engine.TRITON`` but the inputs are not CUDA - float32, or Triton is unavailable. - """ - eng = engine if engine is not None else get_engine() - if eng is Engine.EAGER: - return False - - cuda_f32 = all( - t is None or (t.is_cuda and t.dtype is torch.float32) for t in tensors - ) - - if eng is Engine.TRITON: - if not cuda_f32: - raise RuntimeError( - "engine=Engine.TRITON requires CUDA float32 inputs" - ) - if not triton_available(): - raise RuntimeError("Triton is not available") - return True - - # AUTO - return cuda_f32 and triton_available() diff --git a/torchref/utils/utils.py b/torchref/utils/utils.py index 3e5d0366..06a4f85f 100644 --- a/torchref/utils/utils.py +++ b/torchref/utils/utils.py @@ -257,11 +257,12 @@ class TensorMasks(DeviceMovementMixin, dict): def __init__(self, data=None, device=None): super().__init__() - if device is None: - from torchref.config import get_default_device + from torchref.config import normalize_device - device = get_default_device() - self.device = torch.device(device) + # ``normalize_device`` rather than ``torch.device(...)``: the latter + # keeps an un-indexed spelling ("mps"), which compares unequal to the + # indexed device every real tensor reports. + self.device = normalize_device(device) self._cache = None self._updated = True @@ -314,12 +315,14 @@ def _apply(self, fn): self._cache = None self._updated = True - # Refresh the ``device`` tracker so future ``__setitem__`` calls - # (which migrate incoming tensors to ``self.device``) land correctly. - for v in self.values(): - if isinstance(v, torch.Tensor): - self.device = v.device - break + # Refresh the ``device`` tracker (future ``__setitem__`` calls migrate + # incoming tensors to it) through the shared helper rather than a local + # copy: it also handles the empty-``TensorMasks`` case, where there is + # no mask tensor to read a device from and the tracker has to come from + # the recorded ``.to()`` request. + from torchref.utils.device_mixin import _refresh_device_trackers + + _refresh_device_trackers(self, fn) return self def reset_cache(self) -> None: