RealSpace kernel: coalesced IO + tunable launch; profiler: fused-path timing and --torch_profile

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)
 

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)