RealSpace kernel: coalesced IO + tunable launch; profiler: fused-path timing and --torch_profile - #13
Merged
Merged
Conversation
…ime fused paths The RealSpace Triton kernels loaded/stored each angular column separately at stride n_ang — N_ANG transactions per operand walking DRAM at one-third sector utilization, which capped this bandwidth-bound op well below the A100's bandwidth (measured ~15% at 44k edges x 256 features). Both kernels now do one dense (BLOCK, N_ANG) coalesced tile load/store per operand and pull columns out of the register tile with a multiply-mask reduction (free next to the loads it replaces; no 3D intermediates, so no spill pressure). Launch config is overridable per GPU via ECENET_RS_BLOCK / ECENET_RS_WARPS for a quick sweep. Numerics unchanged: tests/test_realspace_kernel.py passes (CPU paths; run the Triton test on a GPU box). profile_step.py previously timed the *unfused reference* ops for the edge-frame rotate and the MP pack/unrotate/rotate-back/unpack even when the fused flags were on, so the sub-component lines neither summed to the real forward nor showed the fused kernels' cost. The reference lines are now labeled "(unfused ref)" when a fused path replaces them, and the fused ops (edge_frame_fused, pack_unrotate_fused, edge_frame_fused_single) get their own timing lines; nonlin lines are tagged "(fused)" when that path is engaged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
…us copies The launch-config sweep on the A100 showed a ~1.4 ms floor that BLOCK/WARPS could not move: the cost was outside the kernel. EquivariantLinear's einsum emits (n_e, F, n_ang) with strides (F, 1, n_e*F), so the wrapper's .contiguous() was paying a full read+write copy of both inputs every forward call (and of inputs plus both incoming grads in the backward) before the kernel launched. That layout is itself perfectly coalesced (rows are unit-stride), so the copies bought nothing. Each operand now carries its own (row, col) strides into the kernel, and the _rows() helper flattens (e, f) to rows without copying whenever stride(0) == F*stride(1) — covering both dense row-major and the einsum m-major layout, with a contiguous fallback for anything else. GPU test extended with an m-major case (forward and backward) against the contiguous result. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
…r on A100 The previous commit read EquivariantLinear's m-major einsum output in place (strides (F, 1, R)) to skip the wrapper's contiguous copies. Profiling on the A100 showed the opposite of the intent: the fused nonlinearity went from 1.38 ms to ~6.0 ms per instance (512-atom box, 44k edges), with the same regression inside the MP trunk and receiver lines and a proportional hit in the backward — the (BLOCK, N_ANG) tile load degrades to strided scalar accesses across three ~45 MB-apart streams instead of one vectorized coalesced read. _rows() now skips the copy only for genuine row-major layouts (unit-stride coefficient axis); the m-major layout goes back to the contiguous copy the kernel is fast on. Per-operand stride plumbing stays for a future transposed-tile experiment. Restores the prior 159.6 ms joint-step timing. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
The per-section timers only see the forward, op by op; the backward (the larger half of a force call — currently ~61% of the joint step) and the kernel-level cost of composite blocks (receiver, trunk) are invisible to them. --torch_profile reruns the same Totals closures under torch.profiler and prints three aggregate tables: forward-only by self CUDA time, the joint E+F step by self CUDA time (rows that grow between the two are the backward; record_function tags mark the Triton regions), and the joint step by self CUDA memory for the saved-tensor footprint behind large-box OOMs. --profile_rows controls table length. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Performance work on the fused RealSpace nonlinearity plus profiler upgrades, measured on Perlmutter A100 (512-atom diamond box, 44k edges, float32 + TF32, all fusion flags on):
(BLOCK, N_ANG)tile load/store per operand, columns extracted in registers. Per-instance cost 2.32 → 1.38 ms across all 10 module instances; joint E+F step 174.1 → 159.6 ms.ECENET_RS_BLOCK,ECENET_RS_WARPS); a 12-point sweep on the A100 showed the default (128/4) is at the plateau._rows()documents the measurement; per-operand stride plumbing retained.(unfused ref)when a fused kernel replaces them; the fused ops (edge_frame_fused,pack_unrotate_fused,edge_frame_fused_single) get their own lines; nonlin lines tagged(fused).--torch_profile— per-op forward/backward/memory tables via torch.profiler; this is what identified the next optimization target (einsum layout-churn copies, ~29% of the joint step).Test plan
tests/test_realspace_kernel.pypasses on CPU (forward/backward equivalence, finite difference, conservativity, SO(3)); extended with an m-major-layout GPU case.python tests/test_realspace_kernel.pyonce on a CUDA box before merging to exercise thetriton vs fp64 refandtriton m-major layouttests (they skip on CPU).🤖 Generated with Claude Code
https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK