diff --git a/ecenet/realspace_kernel.py b/ecenet/realspace_kernel.py index 09a0599..a5332fb 100644 --- a/ecenet/realspace_kernel.py +++ b/ecenet/realspace_kernel.py @@ -32,6 +32,8 @@ rms-norm / channel-mix variants) from silently taking the wrong path. """ +import os + import torch try: @@ -41,7 +43,11 @@ except ImportError: # CPU-only / no-triton env → PyTorch path _HAS_TRITON = False -_RS_BLOCK = 128 # rows (edge·feature) per program; tunable +# Launch configuration. Rows (edge·feature) per program and warps per program; +# overridable per GPU without a code change (read once at import): +# ECENET_RS_BLOCK=256 ECENET_RS_WARPS=8 python ... +_RS_BLOCK = int(os.environ.get('ECENET_RS_BLOCK', 128)) +_RS_WARPS = int(os.environ.get('ECENET_RS_WARPS', 4)) def is_fusible(nl): @@ -80,6 +86,16 @@ def realspace_reference(A_cos, A_sin, cos_synth, sin_synth, # store. The fat grid tensor never touches HBM, and the backward recomputes it (so # it's never saved) — that's the memory win. # +# Memory access: the op is bandwidth-bound (a handful of flops per loaded float), +# so all global traffic goes through ONE (BLOCK, N_ANG) tile load/store per +# operand — a dense row-major region, fully coalesced. The earlier per-m column +# loads walked DRAM at stride n_ang (a third of each sector wasted) and issued +# N_ANG separate transactions per operand; that access pattern, not arithmetic, +# capped the kernel well below bandwidth. Columns are then pulled out of the +# register tile with a multiply-mask reduction — arithmetically free next to the +# loads it replaces, and it keeps the live set to 2D tiles (no (BLOCK, N_ANG, +# N_GRID) intermediate to spill). +# # n_ang (3-4) and n_grid (9-13) are far below tl.dot's 16-min, so synthesis and # analysis are written as explicit outer-product accumulations / axis reductions # over the (compile-time) angular and grid dims — small 2D tiles only, no tl.dot, @@ -92,51 +108,85 @@ def realspace_reference(A_cos, A_sin, cos_synth, sin_synth, @triton.jit def _rs_fwd_kernel(acos_ptr, asin_ptr, cs_ptr, ss_ptr, ca_ptr, sa_ptr, oc_ptr, os_ptr, - R, n_ang, n_grid, s_row, s_col, + R, n_ang, n_grid, + sc_row, sc_col, ss_row, ss_col, so_row, so_col, BLOCK: tl.constexpr, N_ANG: tl.constexpr, N_GRID: tl.constexpr): offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) # row = edge·feature mask = offs < R + a_idx = tl.arange(0, N_ANG) offs_g = tl.arange(0, N_GRID) g_mask = offs_g < n_grid + # One coalesced (BLOCK, N_ANG) tile load per operand. Each operand + # carries its own (row, col) strides so both supported layouts — dense + # row-major and einsum's m-major (rows unit-stride, see _rows) — read + # straight from where the producer left them, no .contiguous() copies. + t_mask = mask[:, None] & (a_idx < n_ang)[None, :] + ac_all = tl.load(acos_ptr + offs[:, None] * sc_row + a_idx[None, :] * sc_col, + mask=t_mask, other=0.0) + as_all = tl.load(asin_ptr + offs[:, None] * ss_row + a_idx[None, :] * ss_col, + mask=t_mask, other=0.0) + # Synthesis: f(grid) = Σ_m A_cos[:,m]·cos_synth[m,:] + A_sin[:,m]·sin_synth[m,:] + # (column m extracted from the register tile by multiply-mask reduction) f = tl.zeros((BLOCK, N_GRID), dtype=tl.float32) for m in tl.static_range(N_ANG): - col = mask & (m < n_ang) - ac = tl.load(acos_ptr + offs * s_row + m * s_col, mask=col, other=0.0) - as_ = tl.load(asin_ptr + offs * s_row + m * s_col, mask=col, other=0.0) + sel = (a_idx == m).to(tl.float32) + ac = tl.sum(ac_all * sel[None, :], axis=1) + as_ = tl.sum(as_all * sel[None, :], axis=1) cs = tl.load(cs_ptr + m * n_grid + offs_g, mask=g_mask & (m < n_ang), other=0.0) ss = tl.load(ss_ptr + m * n_grid + offs_g, mask=g_mask & (m < n_ang), other=0.0) f += ac[:, None] * cs[None, :] + as_[:, None] * ss[None, :] h = f * tl.sigmoid(f) # silu - # Analysis: out[:,a] = Σ_k h[:,k]·cos_analysis[k,a] (store each column directly) + # Analysis: out[:,a] = Σ_k h[:,k]·cos_analysis[k,a] — accumulated into + # (BLOCK, N_ANG) register tiles, one coalesced store per operand. + oc_all = tl.zeros((BLOCK, N_ANG), dtype=tl.float32) + os_all = tl.zeros((BLOCK, N_ANG), dtype=tl.float32) for a in tl.static_range(N_ANG): col = g_mask & (a < n_ang) ca = tl.load(ca_ptr + offs_g * n_ang + a, mask=col, other=0.0) sa = tl.load(sa_ptr + offs_g * n_ang + a, mask=col, other=0.0) - oc = tl.sum(h * ca[None, :], axis=1) - os = tl.sum(h * sa[None, :], axis=1) - tl.store(oc_ptr + offs * s_row + a * s_col, oc, mask=mask & (a < n_ang)) - tl.store(os_ptr + offs * s_row + a * s_col, os, mask=mask & (a < n_ang)) + sel = (a_idx == a).to(tl.float32) + oc_all += tl.sum(h * ca[None, :], axis=1)[:, None] * sel[None, :] + os_all += tl.sum(h * sa[None, :], axis=1)[:, None] * sel[None, :] + tile_o = offs[:, None] * so_row + a_idx[None, :] * so_col + tl.store(oc_ptr + tile_o, oc_all, mask=t_mask) + tl.store(os_ptr + tile_o, os_all, mask=t_mask) @triton.jit def _rs_bwd_kernel(acos_ptr, asin_ptr, cs_ptr, ss_ptr, ca_ptr, sa_ptr, goc_ptr, gos_ptr, dac_ptr, das_ptr, - R, n_ang, n_grid, s_row, s_col, + R, n_ang, n_grid, + sc_row, sc_col, ss_row, ss_col, + gc_row, gc_col, gs_row, gs_col, so_row, so_col, BLOCK: tl.constexpr, N_ANG: tl.constexpr, N_GRID: tl.constexpr): offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) mask = offs < R + a_idx = tl.arange(0, N_ANG) offs_g = tl.arange(0, N_GRID) g_mask = offs_g < n_grid + # One coalesced (BLOCK, N_ANG) tile load per operand, each with its own + # strides (see forward — inputs and incoming grads have independent + # producers, so their layouts are independent). + t_mask = mask[:, None] & (a_idx < n_ang)[None, :] + ac_all = tl.load(acos_ptr + offs[:, None] * sc_row + a_idx[None, :] * sc_col, + mask=t_mask, other=0.0) + as_all = tl.load(asin_ptr + offs[:, None] * ss_row + a_idx[None, :] * ss_col, + mask=t_mask, other=0.0) + goc_all = tl.load(goc_ptr + offs[:, None] * gc_row + a_idx[None, :] * gc_col, + mask=t_mask, other=0.0) + gos_all = tl.load(gos_ptr + offs[:, None] * gs_row + a_idx[None, :] * gs_col, + mask=t_mask, other=0.0) + # Recompute f (never stored in the forward). f = tl.zeros((BLOCK, N_GRID), dtype=tl.float32) for m in tl.static_range(N_ANG): - col = mask & (m < n_ang) - ac = tl.load(acos_ptr + offs * s_row + m * s_col, mask=col, other=0.0) - as_ = tl.load(asin_ptr + offs * s_row + m * s_col, mask=col, other=0.0) + sel = (a_idx == m).to(tl.float32) + ac = tl.sum(ac_all * sel[None, :], axis=1) + as_ = tl.sum(as_all * sel[None, :], axis=1) cs = tl.load(cs_ptr + m * n_grid + offs_g, mask=g_mask & (m < n_ang), other=0.0) ss = tl.load(ss_ptr + m * n_grid + offs_g, mask=g_mask & (m < n_ang), other=0.0) f += ac[:, None] * cs[None, :] + as_[:, None] * ss[None, :] @@ -146,24 +196,29 @@ def _rs_bwd_kernel(acos_ptr, asin_ptr, cs_ptr, ss_ptr, ca_ptr, sa_ptr, # dh[:,k] = Σ_a g_out_cos[:,a]·cos_analysis[k,a] + g_out_sin[:,a]·sin_analysis[k,a] dh = tl.zeros((BLOCK, N_GRID), dtype=tl.float32) for a in tl.static_range(N_ANG): - col = mask & (a < n_ang) - goc = tl.load(goc_ptr + offs * s_row + a * s_col, mask=col, other=0.0) - gos = tl.load(gos_ptr + offs * s_row + a * s_col, mask=col, other=0.0) + sel = (a_idx == a).to(tl.float32) + goc = tl.sum(goc_all * sel[None, :], axis=1) + gos = tl.sum(gos_all * sel[None, :], axis=1) ca = tl.load(ca_ptr + offs_g * n_ang + a, mask=g_mask & (a < n_ang), other=0.0) sa = tl.load(sa_ptr + offs_g * n_ang + a, mask=g_mask & (a < n_ang), other=0.0) dh += goc[:, None] * ca[None, :] + gos[:, None] * sa[None, :] df = dh * silu_prime # through silu - # dA_cos[:,m] = Σ_k df[:,k]·cos_synth[m,k] + # dA_cos[:,m] = Σ_k df[:,k]·cos_synth[m,k] — accumulated into (BLOCK, + # N_ANG) register tiles, one coalesced store per operand. + dac_all = tl.zeros((BLOCK, N_ANG), dtype=tl.float32) + das_all = tl.zeros((BLOCK, N_ANG), dtype=tl.float32) for m in tl.static_range(N_ANG): col = g_mask & (m < n_ang) cs = tl.load(cs_ptr + m * n_grid + offs_g, mask=col, other=0.0) ss = tl.load(ss_ptr + m * n_grid + offs_g, mask=col, other=0.0) - dac = tl.sum(df * cs[None, :], axis=1) - das = tl.sum(df * ss[None, :], axis=1) - tl.store(dac_ptr + offs * s_row + m * s_col, dac, mask=mask & (m < n_ang)) - tl.store(das_ptr + offs * s_row + m * s_col, das, mask=mask & (m < n_ang)) + sel = (a_idx == m).to(tl.float32) + dac_all += tl.sum(df * cs[None, :], axis=1)[:, None] * sel[None, :] + das_all += tl.sum(df * ss[None, :], axis=1)[:, None] * sel[None, :] + tile_o = offs[:, None] * so_row + a_idx[None, :] * so_col + tl.store(dac_ptr + tile_o, dac_all, mask=t_mask) + tl.store(das_ptr + tile_o, das_all, mask=t_mask) def _can_use_triton(A_cos, activation): @@ -175,21 +230,42 @@ def _fp32(*ts): return [t.to(torch.float32).contiguous() for t in ts] +def _rows(t): + """View a (n_e, F, n_ang) operand as R = n_e·F rows of n_ang coefficients, + returning (tensor, s_row, s_col) for the kernel's stride arithmetic. + + The copy is skipped only for row-major layouts (rows collapse AND the + coefficient axis is unit-stride) — the case the kernel loads as one + vectorized coalesced tile. Reading EquivariantLinear's m-major einsum + output (strides (F, 1, R)) in place was tried and MEASURED ~4-6x slower + on an A100 than copying it contiguous first, despite its unit-stride + rows — the tile load degrades to strided scalar accesses across three + ~45 MB-apart streams. So the m-major layout takes the contiguous copy; + revisit only with a transposed-tile load variant benchmarked in hand.""" + t = t.to(torch.float32) + n_e, F, n_ang = t.shape + if t.stride(2) != 1 or t.stride(0) != F * t.stride(1): + t = t.contiguous() + return t, t.stride(1), t.stride(2) + + def _realspace_forward_triton(A_cos, A_sin, cos_synth, sin_synth, cos_analysis, sin_analysis): n_e, F, n_ang = A_cos.shape n_grid = cos_synth.shape[1] R = n_e * F - acos, asin, cs, ss, ca, sa = _fp32( - A_cos.reshape(R, n_ang), A_sin.reshape(R, n_ang), - cos_synth, sin_synth, cos_analysis, sin_analysis) + acos, sc_r, sc_c = _rows(A_cos) + asin, ss_r, ss_c = _rows(A_sin) + cs, ss, ca, sa = _fp32(cos_synth, sin_synth, cos_analysis, sin_analysis) oc = torch.empty(R, n_ang, device=A_cos.device, dtype=torch.float32) os = torch.empty_like(oc) N_ANG, N_GRID = triton.next_power_of_2(n_ang), triton.next_power_of_2(n_grid) grid = (triton.cdiv(R, _RS_BLOCK),) _rs_fwd_kernel[grid](acos, asin, cs, ss, ca, sa, oc, os, - R, n_ang, n_grid, acos.stride(0), acos.stride(1), - BLOCK=_RS_BLOCK, N_ANG=N_ANG, N_GRID=N_GRID) + R, n_ang, n_grid, + sc_r, sc_c, ss_r, ss_c, oc.stride(0), oc.stride(1), + BLOCK=_RS_BLOCK, N_ANG=N_ANG, N_GRID=N_GRID, + num_warps=_RS_WARPS) out = lambda t: t.reshape(n_e, F, n_ang).to(A_cos.dtype) # noqa: E731 return out(oc), out(os) @@ -199,17 +275,21 @@ def _realspace_backward_triton(g_out_cos, g_out_sin, A_cos, A_sin, n_e, F, n_ang = A_cos.shape n_grid = cos_synth.shape[1] R = n_e * F - acos, asin, cs, ss, ca, sa, goc, gos = _fp32( - A_cos.reshape(R, n_ang), A_sin.reshape(R, n_ang), - cos_synth, sin_synth, cos_analysis, sin_analysis, - g_out_cos.reshape(R, n_ang), g_out_sin.reshape(R, n_ang)) + acos, sc_r, sc_c = _rows(A_cos) + asin, ss_r, ss_c = _rows(A_sin) + goc, gc_r, gc_c = _rows(g_out_cos) + gos, gs_r, gs_c = _rows(g_out_sin) + cs, ss, ca, sa = _fp32(cos_synth, sin_synth, cos_analysis, sin_analysis) dac = torch.empty(R, n_ang, device=A_cos.device, dtype=torch.float32) das = torch.empty_like(dac) N_ANG, N_GRID = triton.next_power_of_2(n_ang), triton.next_power_of_2(n_grid) grid = (triton.cdiv(R, _RS_BLOCK),) _rs_bwd_kernel[grid](acos, asin, cs, ss, ca, sa, goc, gos, dac, das, - R, n_ang, n_grid, acos.stride(0), acos.stride(1), - BLOCK=_RS_BLOCK, N_ANG=N_ANG, N_GRID=N_GRID) + R, n_ang, n_grid, + sc_r, sc_c, ss_r, ss_c, + gc_r, gc_c, gs_r, gs_c, dac.stride(0), dac.stride(1), + BLOCK=_RS_BLOCK, N_ANG=N_ANG, N_GRID=N_GRID, + num_warps=_RS_WARPS) out = lambda t: t.reshape(n_e, F, n_ang).to(A_cos.dtype) # noqa: E731 return out(dac), out(das) diff --git a/tests/test_realspace_kernel.py b/tests/test_realspace_kernel.py index c37df24..2c10fa5 100644 --- a/tests/test_realspace_kernel.py +++ b/tests/test_realspace_kernel.py @@ -197,6 +197,32 @@ def run(fn): print(f" triton vs fp64 ref (F={F}, m_max={m_max}): " f"fwd {f_err:.1e}, dA {b_err:.1e}") + # m-major layout (strides (F, 1, n_e·F) — what EquivariantLinear's + # einsum emits): the kernel reads it via strides, no contiguous copy. + # Must match the contiguous-input result bit-for-bit-ish. + def as_m_major(t): + out = torch.empty(t.shape[2], t.shape[0], t.shape[1], + device=t.device, dtype=t.dtype).permute(1, 2, 0) + out.copy_(t) + return out + + def run_mm(fn): + a = as_m_major(A_cos).requires_grad_(True) + b = as_m_major(A_sin).requires_grad_(True) + oc, os = fn(a, b) + (oc * goc + os * gos).sum().backward() + return oc, os, a.grad, b.grad + + mm = run_mm(lambda a, b: RealSpaceFused.apply( + a, b, nl.cos_synth, nl.sin_synth, nl.cos_analysis, nl.sin_analysis, + nl.activation)) + f_mm = max((ker[0] - mm[0]).abs().max(), (ker[1] - mm[1]).abs().max()).item() + b_mm = max((ker[2] - mm[2]).abs().max(), (ker[3] - mm[3]).abs().max()).item() + assert f_mm < 1e-6, f"m-major fwd F={F},m={m_max}: {f_mm:.2e}" + assert b_mm < 1e-6, f"m-major grad F={F},m={m_max}: {b_mm:.2e}" + print(f" triton m-major layout (F={F}, m_max={m_max}): " + f"fwd {f_mm:.1e}, dA {b_mm:.1e}") + def _structure(n=7, n_types=4, seed=0): g = torch.Generator().manual_seed(seed) diff --git a/tools/profile_step.py b/tools/profile_step.py index 7392e54..2e19e40 100644 --- a/tools/profile_step.py +++ b/tools/profile_step.py @@ -38,6 +38,15 @@ parser.add_argument('--edge_frame_e2n', action='store_true', help="With --edge_frame_fused: ALSO fuse the MP layers' " 'pack+unrotate (their step 3).') +parser.add_argument('--torch_profile', action='store_true', + help='After the timing sections, run the forward and the ' + 'joint E+F step under torch.profiler and print the top ' + 'ops by self CUDA time (plus a top-memory table). This ' + 'is the view the per-section timers cannot give: the ' + 'backward decomposed by op, and record_function-tagged ' + 'kernel regions.') +parser.add_argument('--profile_rows', type=int, default=25, + help='--torch_profile: rows per table (default 25)') args = parser.parse_args() dtype = torch.float32 if args.float32 else torch.float64 @@ -165,23 +174,33 @@ def fresh_pos(): A_emb = model._embed(A, types) A_both = torch.cat([A_emb[edge_i], A_emb[edge_j]], dim=1) -# Wigner rotation breakdown +# Wigner rotation breakdown. When --edge_frame_fused is on, the real forward +# replaces the gather+cat, the bmm rotate, and sph_to_angular with one fused op +# — those lines are then labeled "(unfused ref)" and the fused op is timed too, +# so the sub-components reflect the path that actually runs. +_ef = ' (unfused ref)' if args.edge_frame_fused else '' _, D_list = time_fn(" Wigner D (recursive_wigner_D)", lambda: recursive_wigner_D(r_hat, model.l_max)) _, D_block_main = time_fn(" Wigner D (build_D_block_from_list)", lambda: build_D_block_from_list(D_list, len(r_hat), model.l_max, r_hat.device, r_hat.dtype)) -time_fn(" Wigner bmm (A_both @ D)", +time_fn(f" Wigner bmm (A_both @ D){_ef}", lambda: torch.bmm(A_both, D_block_main)) -time_fn("Wigner rotate total (main, w/ cached D)", +time_fn(f"Wigner rotate total (main, w/ cached D){_ef}", lambda: wigner_rotate(A_both, D_block_main)) A_rot = wigner_rotate(A_both, D_block_main) # sph_to_angular -time_fn("sph_to_angular (repeat_interleave + gather)", +time_fn(f"sph_to_angular (repeat_interleave + gather){_ef}", lambda: model.sph_to_angular(A_rot)) +if args.edge_frame_fused: + from ecenet.edge_frame_kernel import edge_frame_fused, edge_frame_fused_single, pack_unrotate_fused + time_fn("edge-frame fused (gather+rotate+reshape)", + lambda: edge_frame_fused(A_emb, edge_i, edge_j, D_block_main, + model.sph_to_angular)) + A_cos, A_sin = model.sph_to_angular(A_rot) type_i = types[edge_i] @@ -200,12 +219,16 @@ def fresh_pos(): # Low-rank (bottleneck_dim) layers have linear_down/linear_up instead # of a single full-width linear; the nonlinearity runs at the # bottleneck width, between the two. + # The nonlin lines call the real module, so they already time the fused + # path when --fuse_nonlin is on — the suffix just says which path ran. + nl_tag = (' (fused)' if layer.use_nonlinearity + and getattr(layer.nonlin, 'fused', False) else '') if layer.bottleneck_dim is not None: time_fn(f" EquivariantLinear down [{gi},{li}]", lambda l=layer: l.linear_down(A_cos, A_sin)) lin_out = layer.linear_down(A_cos, A_sin) if layer.use_nonlinearity: - time_fn(f" RealSpaceNonlinearity [{gi},{li}]", + time_fn(f" RealSpaceNonlinearity [{gi},{li}]{nl_tag}", lambda l=layer, c=lin_out: l.nonlin(*c)) nl_out = layer.nonlin(*lin_out) if layer.use_nonlinearity else lin_out time_fn(f" EquivariantLinear up [{gi},{li}]", @@ -215,7 +238,7 @@ def fresh_pos(): lambda l=layer: l.linear(A_cos, A_sin)) lin_out = layer.linear(A_cos, A_sin) if layer.use_nonlinearity: - time_fn(f" RealSpaceNonlinearity [{gi},{li}]", + time_fn(f" RealSpaceNonlinearity [{gi},{li}]{nl_tag}", lambda l=layer, c=lin_out: l.nonlin(*c)) A_cos, A_sin = layer(A_cos, A_sin) @@ -236,12 +259,20 @@ def _trunk(): m_sin = u_sin[:, :mp.n_ch] + A_sin s = u_cos[:, mp.n_ch:mp.n_ch + mp.n_scores, 0] - _, h_packed = time_fn(" MP pack cos/sin → n_sph", + # With e2n fusion the real forward runs pack+unrotate as one fused op; + # the two reference lines are kept (labeled) for comparison and the + # fused op is timed as its own line. + _e2n = ' (unfused ref)' if getattr(mp, 'edge_frame_fused_e2n', False) else '' + _, h_packed = time_fn(f" MP pack cos/sin → n_sph{_e2n}", lambda: mp._pack(m_cos, m_sin)) D_block_main_T = D_block_main.transpose(-1, -2) - time_fn(" MP bmm unrotate (h @ D^T)", + time_fn(f" MP bmm unrotate (h @ D^T){_e2n}", lambda: torch.bmm(h_packed, D_block_main_T)) + if getattr(mp, 'edge_frame_fused_e2n', False): + time_fn(" MP pack+unrotate (fused e2n)", + lambda: pack_unrotate_fused(m_cos, m_sin, D_block_main, + mp.l_max, mp.m_max)) h_global = torch.bmm(h_packed, D_block_main_T) def _weights(): @@ -275,10 +306,18 @@ def _weights(): Delta = torch.zeros(n_atoms, H, hb, mp.n_sph, device=device, dtype=dtype ).scatter_add(0, idx, contrib).reshape(n_atoms, mp.n_base, mp.n_sph) - time_fn(" MP bmm rotate back (Delta @ D)", + # Same for the return leg: gather+rotate-back+unpack is one fused op in + # the real forward when edge_frame_fused is on. + _eff = ' (unfused ref)' if getattr(mp, 'edge_frame_fused', False) else '' + time_fn(f" MP bmm rotate back (Delta @ D){_eff}", lambda: torch.bmm(Delta[edge_i], D_block_main)) v_rot = torch.bmm(Delta[edge_i], D_block_main) - _, d = time_fn(" MP unpack n_sph → cos/sin", lambda: mp._unpack(v_rot, n_e)) + _, d = time_fn(f" MP unpack n_sph → cos/sin{_eff}", + lambda: mp._unpack(v_rot, n_e)) + if getattr(mp, 'edge_frame_fused', False): + time_fn(" MP gather+rotate+unpack (fused)", + lambda: edge_frame_fused_single(Delta, edge_i, D_block_main, + mp.l_max, mp.m_max)) time_fn(" MP receiver block", lambda: mp.receiver(*d)) A_cos, A_sin = mp(A_cos, A_sin, r_hat, dist_ij, edge_i, edge_j, @@ -352,3 +391,53 @@ def fwd_forces_les(): time_fn("forward_pbc + E_lr + forces (joint graph)", fwd_forces_les) print() + +# ── 5. torch.profiler breakdown (--torch_profile) ───────────────────────────── +# The per-section timers above can only see the forward, op by op; the backward +# (the larger half of a force call) and the true kernel-level cost of composite +# blocks are invisible to them. This profiles the same closures used in Totals +# and prints per-op aggregates: +# * forward-only table — kernel costs of the real forward path; +# * joint-step table — forward + backward together: rows that grow between +# the two tables (aten::mm vs MmBackward0 etc.) are the backward; the +# record_function tags (edge_frame_fused, pack_unrotate_fused, ...) mark +# our Triton regions in both; +# * joint-step memory table — top self-CUDA allocators, for the saved-tensor +# footprint behind large-box OOMs. +if args.torch_profile: + from torch.profiler import ProfilerActivity, profile + + acts = [ProfilerActivity.CPU] + if device.type == 'cuda': + acts.append(ProfilerActivity.CUDA) + sort_key = ('self_cuda_time_total' if device.type == 'cuda' + else 'self_cpu_time_total') + + def fwd_only(): + return model.forward_pbc(fresh_pos(), types, edge_i, edge_j, + shift_e, nb_src, nb_dst, shift_n) + + joint_fn = fwd_forces_les if les_module is not None else fwd_forces + joint_label = ('joint E_sr + E_lr + forces' if les_module is not None + else 'joint E + forces') + + def run_profile(fn, label, n=args.n_time): + fn() # warm (allocator, autotune caches) + sync() + with profile(activities=acts, profile_memory=True) as prof: + for _ in range(n): + fn() + sync() + ka = prof.key_averages() + sep(f"torch.profiler: {label} (x{n}, top {args.profile_rows} by self " + f"{'CUDA' if device.type == 'cuda' else 'CPU'} time)") + print(ka.table(sort_by=sort_key, row_limit=args.profile_rows)) + return ka + + run_profile(fwd_only, "forward only") + ka_joint = run_profile(joint_fn, joint_label) + + if device.type == 'cuda': + sep(f"torch.profiler: {joint_label} (top 10 by self CUDA memory)") + print(ka_joint.table(sort_by='self_cuda_memory_usage', row_limit=10)) +print()