From e334724c348fd13564b02280acf5ff30beb1d885 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Sat, 10 Oct 2026 15:44:26 +0800 Subject: [PATCH 1/4] perf(pt): assemble Wigner-D behind a custom op under compile Functionalization expands the in-place block scatter into a chain of full-size zeros/index_put adds plus a realized transpose (~4 GB of HBM traffic per frame at the SeZM N=4096 workload). A wigner_d_assemble custom op keeps the eager scatter and hands the graph two contiguous tensors; the traced assembly takes it whenever a fake-tensor mode is active, so the eager path is untouched. --- deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 163 +++++++++++++++--- 1 file changed, 139 insertions(+), 24 deletions(-) diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py index b283e3a7e5..d420dec53d 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py +++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py @@ -45,6 +45,88 @@ _TORCH_RELEASE = Version(torch.__version__).release[:2] +def _wigner_d_assemble_impl( + blocks: list[torch.Tensor], + dim_full: int, + poly_lmin: int, + n_regular: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Assemble the dense block-diagonal matrix and its contiguous transpose. + + 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 -- one buffer, one transpose copy. + """ + 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, d_full.transpose(-1, -2).contiguous() + + +_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, +) -> tuple[torch.Tensor, torch.Tensor]: + ref = blocks[0] + d_full = ref.new_zeros((ref.shape[0], dim_full, dim_full)) + return d_full, d_full.transpose(-1, -2).contiguous() + + +def _wigner_d_assemble_setup( + ctx: Any, inputs: Any, output: Any +) -> None: + blocks, _, poly_lmin, n_regular = inputs + ctx.blocks = blocks + ctx.poly_lmin = poly_lmin + ctx.n_regular = n_regular + + +def _wigner_d_assemble_backward( + ctx: Any, grad_d_full: torch.Tensor, grad_dt_full: torch.Tensor +) -> tuple[list[torch.Tensor], None, None, None]: + # Each block lands in D as itself and in Dt transposed, so its gradient + # is the matching D-slice plus the transpose of the matching Dt-slice. + # The [0, 0] constant receives no gradient. + grads = [] + off = 1 + for block in ctx.blocks[: ctx.n_regular]: + size = block.shape[1] + g = grad_d_full[:, off : off + size, off : off + size] + gt = grad_dt_full[:, off : off + size, off : off + size] + grads.append(g + gt.transpose(-1, -2)) + off += size + if ctx.n_regular < len(ctx.blocks): + po = ctx.poly_lmin * ctx.poly_lmin + g = grad_d_full[:, po:, po:] + gt = grad_dt_full[:, po:, po:] + grads.append(g + gt.transpose(-1, -2)) + return grads, None, None, None + + +_wigner_d_assemble_op.register_autograd( + _wigner_d_assemble_backward, setup_context=_wigner_d_assemble_setup +) + + class CaseCoefficients(nn.Module): """ Polynomial tables for one magnitude-ordered branch of the quaternion Wigner path. @@ -546,62 +628,63 @@ 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 then either + # runs in place (eager) or behind the opaque wigner_d_assemble op + # (compiled). Under tracing, the in-place slice assignments below + # would be functionalized into a chain of full-size zeros/index_put + # adds plus a realized transpose; the custom op keeps the eager + # scatter and hands the graph two plain contiguous tensors. + 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 +699,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,7 +711,39 @@ def forward( lmax=self.lmax, lmin=self.poly_lmin, ) - D_full[:, self.poly_offset :, self.poly_offset :] = D_poly + + # ``is_compiling`` alone is not enough here: this module is traced + # through make_fx fake tensors, where the flag stays False. A fake + # tensor carries its mode, so probe for that instead. + tracing = ( + torch.compiler.is_compiling() + or getattr(blocks[0], "fake_mode", None) is not None + if blocks + else False + ) + if blocks and tracing: + return _wigner_d_assemble_op( + blocks + ([poly_block] if poly_block is not None else []), + self.dim_full, + self.poly_lmin, + len(blocks), + ) + + 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 + off = 1 + for block in blocks: + size = block.shape[1] + D_full[:, off : off + size, off : off + size] = block + off += size + if poly_block is not None: + D_full[:, self.poly_offset :, self.poly_offset :] = poly_block # Consumers address the inverse rotation through explicit strides or # PyTorch strided operators, so the transpose can share D_full's storage. From aea99479029413e4c53ddeed59d0327cec86d763 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Sat, 10 Oct 2026 15:58:33 +0800 Subject: [PATCH 2/4] perf(pt): return the inverse Wigner-D as a view, not a copy The flash and rotation-back kernels already address the rotation through explicit strides, so the contiguous transpose the assembly produced was pure storage: one 660 MB copy per frame that existed only to be read back. The assemble op now returns D_full alone and the caller hands out a transpose view; consumers read it in place and no transpose copy ever exists. --- deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 42 +++++++++---------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py index d420dec53d..1c5e3d9f8c 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py +++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py @@ -50,14 +50,18 @@ def _wigner_d_assemble_impl( dim_full: int, poly_lmin: int, n_regular: int, -) -> tuple[torch.Tensor, torch.Tensor]: - """Assemble the dense block-diagonal matrix and its contiguous transpose. +) -> 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 -- one buffer, one transpose copy. + 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] @@ -71,7 +75,7 @@ def _wigner_d_assemble_impl( if n_regular < len(blocks): po = poly_lmin * poly_lmin d_full[:, po:, po:] = blocks[n_regular] - return d_full, d_full.transpose(-1, -2).contiguous() + return d_full _wigner_d_assemble_op = torch.library.custom_op( @@ -85,15 +89,12 @@ def _( dim_full: int, poly_lmin: int, n_regular: int, -) -> tuple[torch.Tensor, torch.Tensor]: +) -> torch.Tensor: ref = blocks[0] - d_full = ref.new_zeros((ref.shape[0], dim_full, dim_full)) - return d_full, d_full.transpose(-1, -2).contiguous() + return ref.new_zeros((ref.shape[0], dim_full, dim_full)) -def _wigner_d_assemble_setup( - ctx: Any, inputs: Any, output: Any -) -> None: +def _wigner_d_assemble_setup(ctx: Any, inputs: Any, output: Any) -> None: blocks, _, poly_lmin, n_regular = inputs ctx.blocks = blocks ctx.poly_lmin = poly_lmin @@ -101,24 +102,20 @@ def _wigner_d_assemble_setup( def _wigner_d_assemble_backward( - ctx: Any, grad_d_full: torch.Tensor, grad_dt_full: torch.Tensor + ctx: Any, grad_d_full: torch.Tensor ) -> tuple[list[torch.Tensor], None, None, None]: - # Each block lands in D as itself and in Dt transposed, so its gradient - # is the matching D-slice plus the transpose of the matching Dt-slice. - # The [0, 0] constant receives no gradient. + # 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 block in ctx.blocks[: ctx.n_regular]: size = block.shape[1] - g = grad_d_full[:, off : off + size, off : off + size] - gt = grad_dt_full[:, off : off + size, off : off + size] - grads.append(g + gt.transpose(-1, -2)) + grads.append(grad_d_full[:, off : off + size, off : off + size]) off += size if ctx.n_regular < len(ctx.blocks): po = ctx.poly_lmin * ctx.poly_lmin - g = grad_d_full[:, po:, po:] - gt = grad_dt_full[:, po:, po:] - grads.append(g + gt.transpose(-1, -2)) + grads.append(grad_d_full[:, po:, po:]) return grads, None, None, None @@ -722,12 +719,15 @@ def forward( else False ) if blocks and tracing: - return _wigner_d_assemble_op( + d_full = _wigner_d_assemble_op( blocks + ([poly_block] if poly_block is not None else []), self.dim_full, self.poly_lmin, len(blocks), ) + # Stride-addressed consumers read the inverse rotation straight + # from the shared storage; no transpose copy ever exists. + return d_full, d_full.transpose(-1, -2) D_full = torch.zeros( n_edge, From 40fac5908d005fda5fbb90e297e30c170a25dc4d Mon Sep 17 00:00:00 2001 From: longjunhao Date: Thu, 15 Oct 2026 20:34:22 +0800 Subject: [PATCH 3/4] refactor(pt): share the Wigner-D assembly and inverse-layout exit, drop tensor retention --- deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 86 ++++++++++--------- 1 file changed, 46 insertions(+), 40 deletions(-) diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py index 1c5e3d9f8c..f9be1ab451 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py +++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py @@ -96,9 +96,12 @@ def _( def _wigner_d_assemble_setup(ctx: Any, inputs: Any, output: Any) -> None: blocks, _, poly_lmin, n_regular = inputs - ctx.blocks = blocks + # 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 - ctx.n_regular = n_regular def _wigner_d_assemble_backward( @@ -109,11 +112,10 @@ def _wigner_d_assemble_backward( # constant receives no gradient. grads = [] off = 1 - for block in ctx.blocks[: ctx.n_regular]: - size = block.shape[1] + for size in ctx.block_widths: grads.append(grad_d_full[:, off : off + size, off : off + size]) off += size - if ctx.n_regular < len(ctx.blocks): + if ctx.has_poly_tail: po = ctx.poly_lmin * ctx.poly_lmin grads.append(grad_d_full[:, po:, po:]) return grads, None, None, None @@ -626,12 +628,13 @@ def forward( ) n_edge = edge_quaternion.shape[0] - # Degree blocks are computed first; the dense assembly then either - # runs in place (eager) or behind the opaque wigner_d_assemble op - # (compiled). Under tracing, the in-place slice assignments below - # would be functionalized into a chain of full-size zeros/index_put - # adds plus a realized transpose; the custom op keeps the eager - # scatter and hands the graph two plain contiguous tensors. + # 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 @@ -718,41 +721,44 @@ def forward( if blocks else False ) - if blocks and tracing: - d_full = _wigner_d_assemble_op( - blocks + ([poly_block] if poly_block is not None else []), + # 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: + D_full = _wigner_d_assemble_op( + all_blocks, self.dim_full, self.poly_lmin, len(blocks), ) - # Stride-addressed consumers read the inverse rotation straight - # from the shared storage; no transpose copy ever exists. - return d_full, d_full.transpose(-1, -2) - - 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 - off = 1 - for block in blocks: - size = block.shape[1] - D_full[:, off : off + size, off : off + size] = block - off += size - if poly_block is not None: - D_full[:, self.poly_offset :, self.poly_offset :] = poly_block - - # Consumers address the inverse rotation through explicit strides or - # PyTorch strided operators, so the transpose can share D_full's storage. + 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 From e98c691897acd7bd89004eb5633473847775d39d Mon Sep 17 00:00:00 2001 From: longjunhao Date: Thu, 15 Oct 2026 20:34:22 +0800 Subject: [PATCH 4/4] test(pt): focused coverage for the Wigner-D assemble operator --- .../pt/model/test_descriptor_sezm_triton.py | 206 ++++++++++++++++++ 1 file changed, 206 insertions(+) diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index f359f073d6..9094065d64 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -2053,3 +2053,209 @@ 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): + generator = torch.Generator(device="cuda").manual_seed(seed) + return [ + torch.randn(n_edge, w, w, device="cuda", 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="cuda").manual_seed(seed) + w1 = torch.randn( + dense.shape, device="cuda", dtype=dense.dtype, generator=generator + ) + w2 = torch.randn( + dense.shape, device="cuda", dtype=dense.dtype, generator=generator + ) + # cotangent contributions through both D and its transpose + return (dense * w1).sum() + (dense.transpose(-1, -2) * w2).sum() + + @_GPU_KERNELS + def test_op_matches_reference_with_grads(self): + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + _wigner_d_assemble_impl, + _wigner_d_assemble_op, + ) + + 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) + 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)[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, + ) + ) + + @_GPU_KERNELS + def test_opcheck_schema_fake_autograd(self): + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + _wigner_d_assemble_op, + ) + + 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) + args_tail = [] + n_regular = len(blocks) + if poly_lmin is not None: + args_tail = self._make_blocks([4], 8, seed=13) + 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("cuda"): + torch.library.opcheck( + _wigner_d_assemble_op, + (blocks, dim_full, poly_lmin or 11, n_regular), + ) + finally: + torch.set_default_device(saved) + + @_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()