diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index b3cd33798a..f3d2384d91 100644 --- a/deepmd/pt/entrypoints/freeze_pt2.py +++ b/deepmd/pt/entrypoints/freeze_pt2.py @@ -819,15 +819,20 @@ def _export_with_comm_artifact( ) from torch._inductor import config as inductor_config - sample_inputs = _make_comm_sample_inputs(model, nloc=7, device=torch.device("cpu")) - traced = model.forward_common_lower_exportable_with_comm(*sample_inputs) - exported = torch.export.export( - traced, - sample_inputs, - dynamic_shapes=_build_with_comm_dynamic_shapes(sample_inputs), - strict=False, - prefer_deferred_runtime_asserts_over_guards=True, + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + assemble_op_suppressed_for_export, ) + + sample_inputs = _make_comm_sample_inputs(model, nloc=7, device=torch.device("cpu")) + with assemble_op_suppressed_for_export(): + traced = model.forward_common_lower_exportable_with_comm(*sample_inputs) + exported = torch.export.export( + traced, + sample_inputs, + dynamic_shapes=_build_with_comm_dynamic_shapes(sample_inputs), + strict=False, + prefer_deferred_runtime_asserts_over_guards=True, + ) _strip_shape_assertions(exported.graph_module) if target_device.type != "cpu": from torch.export.passes import ( @@ -1038,22 +1043,32 @@ def _freeze_sezm_to_pt2( # Each model's exportable signature matches its sample tuple positionally # (energy / native-spin edge ABI, or virtual-spin nlist ABI), so a single # splat covers all three contracts. - log.info("Tracing the lower graph on CPU (make_fx)...") - traced = model.forward_common_lower_exportable(*sample_inputs_cpu) - - # Output key order is read from the static FX output node. A CUDA-target - # trace may already contain CUDA-only custom operators and therefore must - # not be replayed on the CPU tracing inputs. - output_keys = traced_output_keys(traced) - - log.info("Exporting the traced graph (torch.export)...") - exported = torch.export.export( - traced, - sample_inputs_cpu, - dynamic_shapes=_build_dynamic_shapes(sample_inputs_cpu), - strict=False, - prefer_deferred_runtime_asserts_over_guards=True, + # The Wigner assembly operator cannot ride an AOTInductor archive (the + # C++ package loader resolves custom ops by name and nothing registers + # sezm_nn::wigner_d_assemble there), so the export trace keeps the eager + # scatter the operator wraps. The in-process inference path is + # unaffected. + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + assemble_op_suppressed_for_export, ) + + log.info("Tracing the lower graph on CPU (make_fx)...") + with assemble_op_suppressed_for_export(): + traced = model.forward_common_lower_exportable(*sample_inputs_cpu) + + # Output key order is read from the static FX output node. A + # CUDA-target trace may already contain CUDA-only custom operators + # and therefore must not be replayed on the CPU tracing inputs. + output_keys = traced_output_keys(traced) + + log.info("Exporting the traced graph (torch.export)...") + exported = torch.export.export( + traced, + sample_inputs_cpu, + dynamic_shapes=_build_dynamic_shapes(sample_inputs_cpu), + strict=False, + prefer_deferred_runtime_asserts_over_guards=True, + ) _strip_shape_assertions(exported.graph_module) # move_to_device_pass handles FakeTensor device propagation cleanly; diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py index b283e3a7e5..f0229c1b66 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py +++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py @@ -11,15 +11,24 @@ ) import math +from contextlib import ( + contextmanager, +) from itertools import ( permutations, product, ) from typing import ( + TYPE_CHECKING, Any, ClassVar, ) +if TYPE_CHECKING: + from collections.abc import ( + Iterator, + ) + import torch import torch.nn as nn from packaging.version import ( @@ -44,6 +53,130 @@ _TORCH_RELEASE = Version(torch.__version__).release[:2] +# ``torch.compiler.is_compiling`` is absent before torch 2.3. The eager +# Wigner path must stay runnable on those versions, so resolve the callable +# once at import time instead of calling it unconditionally. +_COMPILER_IS_COMPILING = getattr( + getattr(torch, "compiler", None), + "is_compiling", + None, +) + +# Export traces keep the eager assembly: AOTInductor lowers a plain +# custom_op to a proxy-executor node resolved by name in the C++ runtime, +# where nothing registers ``sezm_nn::wigner_d_assemble`` -- an exported +# .pt2 built through the operator fails to load. The in-process inference +# path lowers in the same process and resolves the operator through the +# dispatcher, so only the freeze/export seam raises this flag. +_assemble_op_suppressed_for_export = False + + +@contextmanager +def assemble_op_suppressed_for_export() -> Iterator[None]: + """Route export traces through the eager scatter instead of the operator.""" + global _assemble_op_suppressed_for_export + saved = _assemble_op_suppressed_for_export + _assemble_op_suppressed_for_export = True + try: + yield + finally: + _assemble_op_suppressed_for_export = saved + + +def _wigner_d_assemble_impl( + blocks: list[torch.Tensor], + dim_full: int, + poly_lmin: int, + n_regular: int, +) -> torch.Tensor: + """Assemble the dense block-diagonal matrix; the transpose stays a view. + + In-place slice assignments are exactly what Inductor's functionalization + cannot keep in place: traced, they expand into a chain of full-size + ``zeros + index_put + add`` tensors plus a realized transpose (~4 GB of + HBM traffic at the SeZM N=4096 workload). Behind a custom op the eager + scatter stays a scatter. Only ``D_full`` is produced here; the caller + makes ``Dt_full`` a plain transpose view outside the op, and every + consumer addresses it through strides (the flash and rotation kernels + take explicit stride arguments; the eager consumers are strided torch + operators). + """ + ref = blocks[0] + n_edge = ref.shape[0] + d_full = torch.zeros(n_edge, dim_full, dim_full, dtype=ref.dtype, device=ref.device) + d_full[:, 0, 0] = 1.0 + off = 1 + for block in blocks[:n_regular]: + size = block.shape[1] + d_full[:, off : off + size, off : off + size] = block + off += size + if n_regular < len(blocks): + po = poly_lmin * poly_lmin + d_full[:, po:, po:] = blocks[n_regular] + return d_full + + +_wigner_d_assemble_op: Any = None +try: + # torch.library.custom_op / register_fake / register_autograd arrived in + # torch 2.4, and the schema inference for ``list[torch.Tensor]`` + # parameters only supports this signature on later releases still + # (torch 2.4-2.6 reject it outright; 2.1-2.3 lack the API). The + # declared minimum is 2.1, so any registration failure keeps the + # autograd-through-scatter implementation for both the eager and the + # traced route below instead of failing the import. + _wigner_d_assemble_op = torch.library.custom_op( + "sezm_nn::wigner_d_assemble", mutates_args=() + )(_wigner_d_assemble_impl) + + @_wigner_d_assemble_op.register_fake + def _( + blocks: list[torch.Tensor], + dim_full: int, + poly_lmin: int, + n_regular: int, + ) -> torch.Tensor: + ref = blocks[0] + return ref.new_zeros((ref.shape[0], dim_full, dim_full)) +except Exception: # pragma: no cover - depends on the installed torch + _wigner_d_assemble_op = None + + +def _wigner_d_assemble_setup(ctx: Any, inputs: Any, output: Any) -> None: + blocks, _, poly_lmin, n_regular = inputs + # Only the layout metadata is needed to slice the cotangent; storing the + # block tensors would extend every activation's lifetime past the + # forward, which the backward never reads. + ctx.block_widths = [int(block.shape[1]) for block in blocks[:n_regular]] + ctx.has_poly_tail = n_regular < len(blocks) + ctx.poly_lmin = poly_lmin + + +def _wigner_d_assemble_backward( + ctx: Any, grad_d_full: torch.Tensor +) -> tuple[list[torch.Tensor], None, None, None]: + # The transpose view's gradient reaches D_full through the ordinary + # autograd of aten.transpose, so only D-slices arrive here. The [0, 0] + # constant receives no gradient. + grads = [] + off = 1 + for size in ctx.block_widths: + grads.append(grad_d_full[:, off : off + size, off : off + size]) + off += size + if ctx.has_poly_tail: + po = ctx.poly_lmin * ctx.poly_lmin + grads.append(grad_d_full[:, po:, po:]) + return grads, None, None, None + + +if _wigner_d_assemble_op is not None: + try: + _wigner_d_assemble_op.register_autograd( + _wigner_d_assemble_backward, setup_context=_wigner_d_assemble_setup + ) + except Exception: # pragma: no cover - depends on the installed torch + _wigner_d_assemble_op = None + class CaseCoefficients(nn.Module): """ @@ -546,62 +679,64 @@ def forward( eps=self.eps, ) n_edge = edge_quaternion.shape[0] - D_full = torch.zeros( - n_edge, - self.dim_full, - self.dim_full, - dtype=edge_quaternion.dtype, - device=edge_quaternion.device, - ) - D_full[:, 0, 0] = 1.0 + + # Degree blocks are computed first; the dense assembly below is one + # shared implementation, reached directly in eager execution and + # behind the opaque wigner_d_assemble operator under tracing (where + # functionalization would otherwise expand the in-place scatter into + # a chain of full-size zeros/index_put adds plus a realized + # transpose). The operator hands the graph one contiguous tensor; + # the inverse rotation leaves through the shared transpose-view exit. + blocks: list[torch.Tensor] = [] + poly_block: torch.Tensor | None = None if self.lmax >= 1: with nvtx_range("WignerD/l1"): - D_full[:, 1:4, 1:4] = self._compute_l1_block(edge_quaternion) + blocks.append(self._compute_l1_block(edge_quaternion)) if self.lmax >= 2: with nvtx_range("WignerD/l2"): - D_full[:, 4:9, 4:9] = self._compute_l2_block(edge_quaternion) + blocks.append(self._compute_l2_block(edge_quaternion)) if self.lmax >= 3: if self.lmax >= 4: with nvtx_range("WignerD/l3l4"): D_l3, D_l4 = self._compute_l3l4_blocks(edge_quaternion) - D_full[:, 9:16, 9:16] = D_l3 - D_full[:, 16:25, 16:25] = D_l4 + blocks.append(D_l3) + blocks.append(D_l4) else: with nvtx_range("WignerD/l3"): - D_full[:, 9:16, 9:16] = self._compute_l3_block(edge_quaternion) + blocks.append(self._compute_l3_block(edge_quaternion)) if self.lmax >= 5: if self.lmax >= 6: with nvtx_range("WignerD/l5l6"): D_l5, D_l6 = self._compute_l5l6_blocks(edge_quaternion) - D_full[:, 25:36, 25:36] = D_l5 - D_full[:, 36:49, 36:49] = D_l6 + blocks.append(D_l5) + blocks.append(D_l6) else: with nvtx_range("WignerD/l5"): - D_full[:, 25:36, 25:36] = self._compute_l5_block(edge_quaternion) + blocks.append(self._compute_l5_block(edge_quaternion)) if self.lmax >= 7: if self.lmax >= 8: with nvtx_range("WignerD/l7l8"): D_l7, D_l8 = self._compute_l7l8_blocks(edge_quaternion) - D_full[:, 49:64, 49:64] = D_l7 - D_full[:, 64:81, 64:81] = D_l8 + blocks.append(D_l7) + blocks.append(D_l8) else: with nvtx_range("WignerD/l7"): - D_full[:, 49:64, 49:64] = self._compute_l7_block(edge_quaternion) + blocks.append(self._compute_l7_block(edge_quaternion)) if self.lmax >= 9: if self.lmax >= 10: with nvtx_range("WignerD/l9l10"): D_l9, D_l10 = self._compute_l9l10_blocks(edge_quaternion) - D_full[:, 81:100, 81:100] = D_l9 - D_full[:, 100:121, 100:121] = D_l10 + blocks.append(D_l9) + blocks.append(D_l10) else: with nvtx_range("WignerD/l9"): - D_full[:, 81:100, 81:100] = self._compute_l9_block(edge_quaternion) + blocks.append(self._compute_l9_block(edge_quaternion)) if self.lmax >= self.poly_lmin: with nvtx_range("WignerD/polynomial"): @@ -616,7 +751,7 @@ def forward( self.poly_coeffs, dtype=self.dtype, ) - D_poly = self._wigner_d_pair_to_real( + poly_block = self._wigner_d_pair_to_real( D_re, D_im, ( @@ -628,16 +763,62 @@ def forward( lmax=self.lmax, lmin=self.poly_lmin, ) - D_full[:, self.poly_offset :, self.poly_offset :] = D_poly - # Consumers address the inverse rotation through explicit strides or - # PyTorch strided operators, so the transpose can share D_full's storage. + # ``is_compiling`` alone is not enough: the production eval trace + # goes through make_fx fake tensors, where the flag stays False, and + # under current dynamo the assembly graph-breaks out of torch.compile + # regions (the flag is False in the fallback too). A fake tensor + # carries its mode, so probe for that as well; the ``is_compiling`` + # disjunct stays for compiled regions that enclose this code. On + # torch older than 2.3 the callable does not exist and the fake-tensor + # probe carries the decision alone. + compiling = ( + _COMPILER_IS_COMPILING() if _COMPILER_IS_COMPILING is not None else False + ) + fake_mode = bool(blocks) and getattr(blocks[0], "fake_mode", None) is not None + tracing = compiling or fake_mode + # One assembly implementation serves both routes: the eager call + # goes straight to the differentiable scatter (autograd flows through + # the slice assignments), the traced call hides it behind the opaque + # operator so functionalization cannot expand the scatter. + all_blocks = blocks + ([poly_block] if poly_block is not None else []) + if not all_blocks: + # lmax == 0: the rotation acts only on the scalar channel. + D_full = torch.zeros( + n_edge, + self.dim_full, + self.dim_full, + dtype=edge_quaternion.dtype, + device=edge_quaternion.device, + ) + D_full[:, 0, 0] = 1.0 + elif ( + tracing + and _wigner_d_assemble_op is not None + and not _assemble_op_suppressed_for_export + ): + D_full = _wigner_d_assemble_op( + all_blocks, + self.dim_full, + self.poly_lmin, + len(blocks), + ) + else: + D_full = _wigner_d_assemble_impl( + all_blocks, + self.dim_full, + self.poly_lmin, + len(blocks), + ) + + # One inverse-layout exit shared by both routes: the compatibility + # policy below is preserved as-is instead of being bypassed by the + # traced branch. Stride-addressed consumers read the inverse + # rotation from the shared storage; PyTorch 2.11 Inductor cannot + # safely lower that escaping transpose view, so it is materialized + # there exactly as the eager route does. Dt_full = D_full.transpose(-1, -2) if self._materialize_inverse_rotation: - # PyTorch 2.11 Inductor cannot safely lower this escaping transpose - # view after the slice assignments that assemble D_full. The - # materialized layout keeps the compiled graph semantically - # identical; later releases retain the shared-storage view. Dt_full = Dt_full.contiguous() return D_full, Dt_full diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index f359f073d6..94be7ea476 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -30,10 +30,14 @@ ) import torch +from packaging.version import ( + Version, +) from torch.fx.experimental.proxy_tensor import ( make_fx, ) +import deepmd.pt.model.descriptor.sezm_nn.wignerd as wignerd_mod from deepmd.pt.model.descriptor.sezm_nn.indexing import ( build_m_major_index, get_so3_dim_of_lmax, @@ -60,6 +64,13 @@ _CUDA = torch.cuda.is_available() +# The assembly operator registers only where the running torch accepts its +# schema (torch >= 2.7); torch 2.4-2.6 have ``torch.library.custom_op`` but +# reject a ``list[torch.Tensor]`` parameter, so tests that call the operator +# directly gate on registration success, not on API presence. +_WIGNER_ASSEMBLE_OP = wignerd_mod._wigner_d_assemble_op +_WIGNER_OP_REGISTERED = _WIGNER_ASSEMBLE_OP is not None + # All SeZM Triton kernels ship together (every per-module availability flag # reduces to "the triton package imports"), so one gate covers every # GPU-kernel test class. Dispatch- and parsing-only classes stay ungated. @@ -2053,3 +2064,463 @@ def test_tune_missing_configs_runs_the_real_flash_bwd_sweep(self): if __name__ == "__main__": unittest.main() + + +class TestSeZMWignerDAssembleOp(unittest.TestCase): + """Focused coverage for the op-bound dense Wigner-D assembly.""" + + WIDTHS = (3, 5, 7) # lmax = 3 ladder + + def _make_blocks(self, widths, n_edge, seed, dtype=torch.float64, device="cuda"): + generator = torch.Generator(device=device).manual_seed(seed) + return [ + torch.randn(n_edge, w, w, device=device, dtype=dtype, generator=generator) + for w in widths + ] + + def _reference_assemble(self, blocks, poly, dim_full, poly_lmin): + ref = blocks[0] + dense = ref.new_zeros(ref.shape[0], dim_full, dim_full) + dense[:, 0, 0] = 1.0 + off = 1 + for block in blocks: + size = block.shape[1] + dense[:, off : off + size, off : off + size] = block + off += size + if poly is not None: + po = poly_lmin * poly_lmin + dense[:, po:, po:] = poly + return dense + + def _loss(self, dense, seed): + generator = torch.Generator(device=dense.device).manual_seed(seed) + w1 = torch.randn( + dense.shape, device=dense.device, dtype=dense.dtype, generator=generator + ) + w2 = torch.randn( + dense.shape, device=dense.device, dtype=dense.dtype, generator=generator + ) + # cotangent contributions through both D and its transpose + return (dense * w1).sum() + (dense.transpose(-1, -2) * w2).sum() + + @unittest.skipUnless( + _WIGNER_OP_REGISTERED, + "the operator registration needs a torch that accepts its schema", + ) + def test_op_matches_reference_with_grads(self): + """Pure-PyTorch operator math; runs wherever the CPU CI runs.""" + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + _wigner_d_assemble_impl, + _wigner_d_assemble_op, + ) + + device = "cuda" if _CUDA else "cpu" + n_edge = 64 + for widths, poly_lmin in ((self.WIDTHS, None), ((3, 5), 3)): + with self.subTest(widths=widths, tail=poly_lmin is not None): + blocks = self._make_blocks(widths, n_edge, seed=5, device=device) + poly = None + n_regular = len(blocks) + if poly_lmin is not None: + poly_w = 4 + poly = self._make_blocks([poly_w], n_edge, seed=7, device=device)[0] + dim_full = sum(widths) + 1 if poly is None else poly_lmin**2 + poly_w + args = ( + blocks + ([poly] if poly is not None else []), + dim_full, + poly_lmin or 11, + n_regular, + ) + + leaves_op = [ + b.clone().requires_grad_(True) + for b in (blocks + ([poly] if poly is not None else [])) + ] + dense_op = _wigner_d_assemble_op(leaves_op, args[1], args[2], args[3]) + leaves_ref = [ + b.clone().requires_grad_(True) + for b in (blocks + ([poly] if poly is not None else [])) + ] + dense_ref = _wigner_d_assemble_impl( + leaves_ref, args[1], args[2], args[3] + ) + self.assertTrue(torch.equal(dense_op.detach(), dense_ref.detach())) + self.assertTrue( + torch.equal( + dense_op.detach(), + self._reference_assemble(blocks, poly, args[1], args[2]), + ) + ) + + grads_op = torch.autograd.grad( + self._loss(dense_op, 11), leaves_op, allow_unused=True + ) + grads_ref = torch.autograd.grad( + self._loss(dense_ref, 11), leaves_ref, allow_unused=True + ) + for g1, g2 in zip(grads_op, grads_ref, strict=True): + torch.testing.assert_close(g1, g2, atol=0.0, rtol=0.0) + + # fp64 gradcheck on the operator route (small edge count) + small = [ + b[:3].clone().detach().requires_grad_(True) + for b in (blocks + ([poly] if poly is not None else [])) + ] + self.assertTrue( + torch.autograd.gradcheck( + lambda *bs: _wigner_d_assemble_op( + list(bs), args[1], args[2], args[3] + ), + small, + ) + ) + + @unittest.skipUnless( + _WIGNER_OP_REGISTERED, + "the operator registration needs a torch that accepts its schema", + ) + def test_opcheck_schema_fake_autograd(self): + """Schema/fake/autograd contract; pure PyTorch, CPU CI included.""" + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + _wigner_d_assemble_op, + ) + + device = "cuda" if _CUDA else "cpu" + for widths, poly_lmin in ((self.WIDTHS, None), ((3, 5), 3)): + with self.subTest(widths=widths, tail=poly_lmin is not None): + blocks = self._make_blocks(widths, 8, seed=9, device=device) + args_tail = [] + n_regular = len(blocks) + if poly_lmin is not None: + args_tail = self._make_blocks([4], 8, seed=13, device=device) + blocks = blocks + args_tail + dim_full = sum(widths) + 1 if not args_tail else poly_lmin**2 + 4 + saved = torch.get_default_device() + torch.set_default_device(None) + try: + with torch.device(device): + torch.library.opcheck( + _wigner_d_assemble_op, + (blocks, dim_full, poly_lmin or 11, n_regular), + ) + finally: + torch.set_default_device(saved) + + @unittest.skipUnless( + _WIGNER_OP_REGISTERED, + "the operator registration needs a torch that accepts its schema", + ) + def test_make_fx_trace_selects_the_assemble_op(self): + """The fake-mode probe half: make_fx tracing picks the opaque op. + + The production eval path lowers through make_fx, whose fake + tensors carry the mode this check exercises (``is_compiling`` + stays False there; under current dynamo the assembly graph-breaks + out of torch.compile regions, so the fake-mode disjunct is the + live one today). The graph-node assertion locks in the mechanism: + with the probe silently false, the traced route would fall back to + the raw scatter and every value comparison would still pass. + """ + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + WignerDCalculator, + ) + + def _run(quat): + return calc.forward(quat) + + calc = WignerDCalculator(lmax=3, dtype=torch.float32).cpu() + calc.eval() + quat = torch.randn(8, 4, dtype=torch.float32, device="cpu") + quat = quat / quat.norm(dim=-1, keepdim=True) + graph = make_fx(_run, tracing_mode="symbolic", _allow_non_fake_inputs=True)( + quat + ) + self.assertIn("sezm_nn.wigner_d_assemble", graph.code) + + @unittest.skipIf( + Version(torch.__version__).release < (2, 7), + "the schema is only accepted from torch 2.7", + ) + def test_operator_registers_on_schema_capable_torch(self): + """A registration regression on a capable torch must fail, not skip. + + The operator tests skip on ``_WIGNER_OP_REGISTERED``, so a regression + that silently turns the flag off would leave the suite green with + the traced route reverted to the scatter. This check is gated on + the torch release instead of the observed flag and therefore fails + when registration breaks on a schema-capable torch (CI pins 2.10 + and 2.11). + """ + self.assertIsNotNone( + wignerd_mod._wigner_d_assemble_op, + "wigner_d_assemble failed to register on a schema-capable torch", + ) + + def test_export_suppression_keeps_the_eager_assembly(self): + """Export traces keep the scatter: AOTI archives cannot carry the op. + + AOTInductor lowers a plain custom_op to a proxy-executor node + resolved by name in the C++ runtime, where nothing registers + ``sezm_nn::wigner_d_assemble`` -- an exported .pt2 built through + the operator fails to load. Under the freeze seam's suppression + context the traced graph must carry the eager scatter instead. + """ + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + assemble_op_suppressed_for_export, + ) + + def _run(quat): + return calc.forward(quat) + + calc = wignerd_mod.WignerDCalculator(lmax=3, dtype=torch.float32).cpu() + calc.eval() + quat = torch.randn(8, 4, dtype=torch.float32, device="cpu") + quat = quat / quat.norm(dim=-1, keepdim=True) + + plain = make_fx(_run, tracing_mode="symbolic", _allow_non_fake_inputs=True)( + quat + ) + self.assertIn("sezm_nn.wigner_d_assemble", plain.code) + + with assemble_op_suppressed_for_export(): + suppressed = make_fx( + _run, tracing_mode="symbolic", _allow_non_fake_inputs=True + )(quat) + self.assertNotIn("sezm_nn.wigner_d_assemble", suppressed.code) + torch.testing.assert_close(suppressed(quat), plain(quat), atol=0.0, rtol=0.0) + + @unittest.skipUnless( + _WIGNER_OP_REGISTERED, + "the operator registration needs a torch that accepts its schema", + ) + def test_is_compiling_half_of_the_probe_selects_the_operator(self): + """The ``is_compiling`` disjunct of the probe must stay live. + + The fake-mode half has its own graph-node pin; this check drives + the other disjunct directly -- a truthy ``is_compiling`` against + real (non-fake) tensors routes the assembly through the operator, + which a dropped disjunct or a broken callable lookup silently + reverts to the scatter. (Under the *current* dynamo the assembly + sits in a graph-break region and torch.compile leaves the flag + False there, so the compiled route itself does not reach this + disjunct yet; the mock pins the mechanism, not dynamo's region + placement.) + """ + calls = [] + real_op = wignerd_mod._wigner_d_assemble_op + + def _spy(*args, **kwargs): + calls.append(1) + return real_op(*args, **kwargs) + + calc = wignerd_mod.WignerDCalculator(lmax=3, dtype=torch.float32).cpu() + calc.eval() + quat = torch.randn(8, 4, dtype=torch.float32, device="cpu") + quat = quat / quat.norm(dim=-1, keepdim=True) + + with ( + mock.patch.object(wignerd_mod, "_COMPILER_IS_COMPILING", lambda: True), + mock.patch.object(wignerd_mod, "_wigner_d_assemble_op", _spy), + ): + out = calc.forward(quat) + self.assertTrue( + bool(calls), + "a truthy is_compiling must route the assembly through the operator", + ) + + def test_import_and_tracing_fall_back_when_registration_fails(self): + """Unsupported torches must import the module and stay eager. + + Two emulations run in subprocesses, one per incompatible class: + ``hide`` removes ``torch.library.custom_op`` for the module's + re-execution (torch 2.1-2.3, where the API is absent -- on a torch + that really lacks it the attribute is already gone and nothing is + deleted), and ``reject`` replaces it with a stub raising the + schema-inference error (torch 2.4-2.6, where the API exists but + rejects a ``list[torch.Tensor]`` parameter). In both, torch's own + lazily importing users of the API stay intact (the change is scoped + to the reload window), the operator stays unregistered, the eager + forward/backward works, and make_fx tracing keeps the eager + implementation instead of the operator node. + + The child first imports the module untouched and asserts the + registration state its parent observed (passed as an argument), so + the successful-registration control is measured against reality + instead of API presence: on a torch whose registration legitimately + fails at import, the control asserts the import-time fallback, and + the operator-specific tests above are skipped by their decorator. + """ + import pathlib + import subprocess + import sys + import textwrap + + tree = pathlib.Path(wignerd_mod.__file__).resolve().parents[5] + child = textwrap.dedent( + """ + import pathlib + import sys + + mode = sys.argv[1] + tree = pathlib.Path(sys.argv[2]) + expected_registered = sys.argv[3] == "1" + sys.meta_path = [ + f + for f in sys.meta_path + if "editable" + not in (type(f).__module__ + "." + type(f).__name__).lower() + and "editable" not in getattr(f, "__name__", "").lower() + ] + sys.path.insert(0, str(tree)) + import importlib + + import torch + import torch.distributed._functional_collectives # noqa: F401 + + import deepmd.pt.model.descriptor.sezm_nn.wignerd as wd + + if expected_registered: + assert wd._wigner_d_assemble_op is not None, "control: registers" + else: + assert wd._wigner_d_assemble_op is None, "control: import fallback" + + saved_custom_op = getattr(torch.library, "custom_op", None) + if mode == "hide": + if saved_custom_op is not None: + del torch.library.custom_op + elif mode == "reject": + def _rejected(*args, **kwargs): + raise TypeError( + "cannot infer schema for blocks: list[torch.Tensor]" + ) + + torch.library.custom_op = _rejected + else: + raise AssertionError(f"unknown mode {mode}") + wd = importlib.reload(wd) + if saved_custom_op is not None: + torch.library.custom_op = saved_custom_op + elif mode == "reject": + # The stub was injected over a missing attribute; remove it. + del torch.library.custom_op + assert wd._wigner_d_assemble_op is None, "fallback: op stays out" + + calc = wd.WignerDCalculator(lmax=1, dtype=torch.float64).cpu() + calc.eval() + quat = torch.randn(4, 4, dtype=torch.float64, device="cpu") + quat = (quat / quat.norm(dim=-1, keepdim=True)).requires_grad_(True) + dense, _dt = calc.forward(quat) + dense.sum().backward() + + from torch.fx.experimental.proxy_tensor import make_fx + + def _run(q): + return calc.forward(q) + + graph = make_fx( + _run, tracing_mode="symbolic", _allow_non_fake_inputs=True + )(quat) + assert ( + "sezm_nn.wigner_d_assemble" not in graph.code + ), "without a registration the traced route keeps the eager path" + print("fallback-ok") + """ + ) + registered = wignerd_mod._wigner_d_assemble_op is not None + for mode in ("hide", "reject"): + with self.subTest(emulation=mode, registered=registered): + result = subprocess.run( + [ + sys.executable, + "-c", + child, + mode, + str(tree), + "1" if registered else "0", + ], + capture_output=True, + text=True, + timeout=300, + ) + self.assertEqual( + result.returncode, + 0, + f"fallback child ({mode}) failed:" + f"\n{result.stdout}\n{result.stderr}", + ) + self.assertIn("fallback-ok", result.stdout) + + @_GPU_KERNELS + def test_traced_matches_eager_and_layout_policy(self): + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + WignerDCalculator, + ) + + generator = torch.Generator(device="cuda").manual_seed(21) + quat = torch.randn(32, 4, device="cuda", generator=generator) + quat = quat / quat.norm(dim=-1, keepdim=True) + calc = WignerDCalculator(lmax=3, dtype=torch.float32).cuda() + calc.eval() + + saved_device = torch.get_default_device() + torch.set_default_device(None) + try: + with torch.device("cuda"): + d_eager, dt_eager = calc.forward(quat) + d_compiled, dt_compiled = torch.compile(calc.forward)(quat) + # The assembly op itself is bitwise (see the reference test); + # the compiled ladder reassociates the degree-block math, so the + # cross-route comparison holds at rounding level. + torch.testing.assert_close(d_compiled, d_eager, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(dt_compiled, dt_eager, atol=1e-5, rtol=1e-5) + finally: + torch.set_default_device(saved_device) + + # The inverse-layout policy is one shared exit: both routes honor + # the compatibility flag instead of the traced route bypassing it. + def _run_routes(): + saved = torch.get_default_device() + torch.set_default_device(None) + try: + with torch.device("cuda"): + _, dt_eager_i = calc.forward(quat) + _, dt_compiled_i = torch.compile(calc.forward)(quat) + return dt_eager_i, dt_compiled_i + finally: + torch.set_default_device(saved) + + calc._materialize_inverse_rotation = False + dt_view, dt_view_c = _run_routes() + self.assertFalse(dt_view.is_contiguous()) + self.assertFalse(dt_view_c.is_contiguous()) + calc._materialize_inverse_rotation = True + dt_dense, dt_dense_c = _run_routes() + self.assertTrue(dt_dense.is_contiguous()) + self.assertTrue(dt_dense_c.is_contiguous()) + torch.testing.assert_close(dt_dense, dt_view, atol=0.0, rtol=0.0) + + @_GPU_KERNELS + def test_lmax_zero_and_empty_edges(self): + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + WignerDCalculator, + ) + + quat = torch.randn(8, 4, device="cuda") + quat = quat / quat.norm(dim=-1, keepdim=True) + calc0 = WignerDCalculator(lmax=0, dtype=torch.float32).cuda() + calc0.eval() + d0, dt0 = calc0.forward(quat) + expected = torch.zeros(8, 1, 1, device="cuda") + expected[:, 0, 0] = 1.0 + self.assertTrue(torch.equal(d0, expected)) + self.assertTrue(torch.equal(dt0, expected)) + + calc3 = WignerDCalculator(lmax=3, dtype=torch.float32).cuda() + calc3.eval() + d_empty, dt_empty = calc3.forward(torch.empty(0, 4, device="cuda")) + self.assertEqual(tuple(d_empty.shape), (0, 16, 16)) + self.assertEqual(tuple(dt_empty.shape), (0, 16, 16)) + + +if __name__ == "__main__": + unittest.main()