perf(pt): assemble Wigner-D behind a custom op under compile

deepmd/pt/model/descriptor/sezm_nn/wignerd.py

@@ -44,6 +44,110 @@
 
 _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,
+)
+
+
+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 +650,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 +722,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 +734,55 @@ 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 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.  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:
+            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