Repository navigation
perf(pt): assemble Wigner-D behind a custom op; inverse rotation as a view - #6070
long-yi-2019 wants to merge 4 commits into
Conversation
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.
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.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 15.38% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 13 functions across 2 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the main changes: custom-op Wigner-D assembly and inverse-rotation handling through a view. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request.
- Fix all pre-merge checks with AI
✨ Finishing Touches: 💡 1
-
🧪 Generate unit tests (beta)
-
- Create a new PR
🛠️ Fix failing CI checks: 💡
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.
Comment @coderabbitai help to get the list of available commands.
|
@OutisLi @wanghan-iapcm @njzjz requesting a review, please. Following the 2026-10-13 discussion we validated the combined stack of our six performance PRs (#6056, #6057, #6060, #6070, #6071, #6072) with a multi-scale gate: diamond + bcc-Cu at ~500 / ~2000 / ~4400 atoms, three columns (pure torch.compile with all custom kernels off / current master / our changes on). All accuracy cells sit within 2.2x the pure baseline's own run-to-run band (criterion: 3x), and no tested atom count regresses in time or memory — full numbers in the PR descriptions. Suggested review order for the six: #6056 -> #6057 -> #6060 (small and self-contained) -> #6070 -> #6071 -> #6072. |
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Reviewed at aea9947.
Purpose and implementation
This PR reduces the temporary dense-tensor work caused by functionalizing Wigner-D block assembly. It collects the per-degree blocks, performs the scatter behind a custom operator during tracing, and returns the inverse rotation as a transpose view. The backward extracts the corresponding block slices. This is a reasonable optimization boundary; it does not change the rotation formulas.
Requested changes
Please remove the unnecessary tensor retention in the autograd context and consolidate the duplicated assembly/return logic, as detailed inline. The latter should make the inverse-layout compatibility policy explicit and shared rather than bypassing it through an early return. Please also update the comment that still promises two contiguous outputs.
The new operator and dispatch have no dedicated tests added by this PR. Please add focused regression coverage for the behavior being introduced, including its gradients and layout/compatibility policy. I have not reproduced a new Wigner-D numerical or gradient defect in the sampled cases; the PyTorch 2.11 concern below is an unverified compatibility change, not a reproduced runtime failure.
Correct higher-order validation
I am excluding the second-order mismatch from my exploratory raw-FX probe from the findings. That probe traced forward plus a first backward, then differentiated the resulting GraphModule again, without the repository's production graph transformations. PyTorch explicitly does not guarantee that construction: maintainer explanation. It is not evidence that this PR or the baseline training implementation has incorrect second derivatives.
For validation:
- Check the operator schema/fake/autograd registration with
torch.library.opcheck, and check the mathematical backward withgradcheck/gradgradcheckor an independent block-assembly reference. Use nontrivial inputs and contributions from both D and its transpose. - For standalone higher-order compilation tests, put the forward, first backward, and required second backward inside the function being traced, then compile that function. Do not assume a raw forward-plus-backward FX graph remains correctly differentiable afterwards.
- For force-loss training parity, exercise the actual model compilation entry point and its existing transformations, including
strip_saved_tensor_detach; compare force-loss gradients for all model parameters against eager execution. Do not replace that pipeline by indiscriminately deleting detach operations. - Cover the zero-degree/empty-edge cases, regular blocks, the polynomial tail, both eager and traced routes, and the intended inverse-layout behavior on supported PyTorch versions. Validate PyTorch 2.11 before removing or bypassing its existing policy.
Checks performed
On this PR head, the seven existing Wigner-D tests pass on CPU. Additional CPU/CUDA checks with FP32/FP64, lmax 0/1/3/6/10/11, empty edges, and the actual custom-op dispatch passed the sampled value/first-derivative/HVP comparisons and opcheck. Ruff, the read-only formatting check, and diff checks pass.
Separately, adding the existing project transformation to the baseline raw-FX probe restores its HVP agreement with eager and with convergent finite differences. The real SeZM test_forward_backward_double_backward_matches_compile also passes on checkout d3b3b66, including force-loss parameter gradients and three optimizer steps. That is a production-path control, not a claim that this PR's complete model suite or every supported PyTorch version was tested. The reported full-model performance gain has not been independently reproduced.
|
|
||
| def _wigner_d_assemble_setup(ctx: Any, inputs: Any, output: Any) -> None: | ||
| blocks, _, poly_lmin, n_regular = inputs | ||
| ctx.blocks = blocks |
There was a problem hiding this comment.
Please store only the block layout metadata in this context, not the block tensors.
The backward reads each block's width and whether the polynomial tail is present; it never needs the tensor values. ctx.blocks = blocks therefore keeps every activation buffer alive unnecessarily. In a CPU and CUDA weak-reference check with three non-leaf blocks, the direct assembly implementation retained none after forward, while the custom op retained all three even after backward; they were released only when the output was released. This establishes an avoidable lifetime extension, not a measured regression in the full compiled-model benchmark.
A tuple of block widths (or slice bounds), together with the tail metadata, is sufficient. Saving the tensors through a different API would still retain data the backward does not need.
| ) | ||
| # 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) |
There was a problem hiding this comment.
Please share the assembly implementation and inverse-layout policy between eager and traced execution.
The block-scatter loop here is duplicated in _wigner_d_assemble_impl, and this early return bypasses _materialize_inverse_rotation at the common eager exit. With that flag enabled, eager returns a contiguous inverse but the traced route returns a non-contiguous view. I verified that branch difference on PyTorch 2.13; I have not run PyTorch 2.11 and am not claiming a reproduced 2.11 failure. The opaque operator may eliminate the original compiler problem, but that needs evidence rather than an implicit policy change.
A cleaner structure is one assembly implementation (invoked directly or through the custom op) followed by one inverse-layout/return step, with the genuine lmax=0 case handled explicitly. Either preserve the existing compatibility policy there, or validate the supported 2.11 path and deliberately revise it with a regression test. Please also fix the comment above that still says the graph receives two contiguous tensors.
|
Thank you for the exceptionally thorough review — and for independently running the value/derivative/HVP comparisons and All three requested changes are addressed in e98c691:
New focused test class Validation at the new head: the full On the PyTorch 2.11 point: agreed that the operator removes the original reason the flag exists on the traced route (the functionalized slice-assignment graph), but since we cannot verify 2.11 here either, the policy is now preserved verbatim in both routes rather than changed implicitly. |
Motivation
Each eval frame assembles the block-diagonal Wigner-D rotation matrices by
scattering the per-degree blocks into a dense
(E, D, D)tensor. Tracedthrough Inductor, those in-place slice assignments functionalize into a
chain of full-size
zeros + index_put + addtensors plus a realizedtranspose -- five
(E, D, D)surfaces combined by one kernel, ~4 GB of HBMtraffic per frame at the SeZM N=4096 workload (E ~ 646k, D = 16). On top
of that, the inverse rotation
Dt_fullwas produced as a full contiguouscopy even though every consumer can address it through strides.
Change
Two commits, one battle:
Assembly behind a custom op.
sezm_nn::wigner_d_assembletakes theper-degree blocks (plus the polynomial tail when
lmax >= 11) and doesthe eager scatter: one zeros buffer, in-place block writes. The traced
forward takes the op whenever a fake-tensor mode is active (
is_compiling()is not enough -- this module is traced through make_fx, where the flag
stays False; a fake tensor carries its mode, so the forward probes for
it). The eager path is untouched. An exact autograd formula returns
each block its D-slice.
The inverse rotation is a view. The flash and rotation-back kernels
already address the rotation through explicit stride arguments
(
_rotation_strides/w_se, w_sr, w_sk), verified bit-for-bit againstthe contiguous copy with non-symmetric matrices. The op now returns
D_fullalone and the caller hands outD_full.transpose(-1, -2); thetranspose copy never exists. The one consumer that requires contiguous
storage (
so2_rotate_mix) readsD_fullitself, which stays contiguous.Measurements
SeZM model, N=4096 atoms (~646k edges), RTX 5090, torch 2.12.1+cu130,
DP_COMPILE_INFER=1,DP_TRITON_INFER=3, throughput mode (16 evals, P50):The step-1 saving is time (the ~4 GB/frame add-chain disappears; the
compiler-fragment family drops 38.3 -> 25.8 ms of kernel time). The step-2
saving is VRAM: deleting the 660 MB transpose copy costs nothing overall
because the stride-read overhead and the removed copy cancel out. Two-state
compiled-eval accuracy at N=4096: energy rel diff 5.8e-10, max force diff
4.8e-6 with max |f| ~ 3.9.
Tests
source/tests/pt/model/test_descriptor_sezm.py: 61 passed (default config).source/tests/pt/model/test_descriptor_sezm_triton.py: 37 passed, 1 skipped (pre-existing).source/tests/pt/model/test_sezm_model.py::TestSeZMModelCompile: 9 passed on this branch head (RTX 5090, torch 2.12.1), including the eval cache/shape-change, first-frame and force-loss double-backward tests.(view vs materialized inverse) matches exactly.
Summary by CodeRabbit