EquivariantLinear: one row-major block-diagonal GEMM — no layout churn

ecenet/equivariant.py

@@ -26,6 +26,27 @@ class EquivariantLinear(nn.Module):
 
     Same weights for cos/sin parts. Bias only on m=0 (invariant).
     Angular channels: m = 0, 1, ..., m_max (index 0 is m=0).
+
+    The map y[e,o,m] = Σ_i x[e,i,m]·W[m,o,i] never mixes angular modes, so it
+    is evaluated as ONE dense row-major GEMM over the flattened (feature,
+    angular) axis, against a (in·n_ang, out·n_ang) weight that is
+    block-diagonal per m (assembled from the parameter each call — ~MB-scale,
+    autograd flows through the assembly; the zero blocks cost n_ang× the
+    strictly needed flops, which TF32/tensor cores absorb). The earlier
+    einsum formulation lowered to a batched bmm over the angular axis, whose
+    batch-contiguity requirement copied the full per-edge activation on every
+    call and emitted an m-major output that downstream consumers (the fused
+    nonlinearity, residual adds, the next einsum) paid to re-normalize —
+    profiling showed that layout churn at ~29% of the whole force step. The
+    dense form's input and output are both plain row-major views: no
+    permutes, no copies, and the layout every consumer wants.
+
+    The trade only pays where the padded GEMM is (near-)free: on CUDA tensor
+    cores. On fp32 CUDA cores (TF32 off) and in float64 the 3x padding costs
+    exactly 3x — measured slower than the einsum path despite the copies —
+    so dispatch is automatic: dense for CUDA fp16/bf16, and for CUDA fp32
+    with TF32 enabled; einsum otherwise. ``dense_gemm`` (True/False)
+    overrides the automatic choice.
     """
 
     def __init__(self, in_features, out_features, n_angular, m_max):
@@ -34,21 +55,52 @@ def __init__(self, in_features, out_features, n_angular, m_max):
         self.out_features = out_features
         self.n_angular = n_angular
         self.m_max = m_max
+        self.dense_gemm = None       # None = auto (see _use_dense)
 
         # (n_angular, out_features, in_features)
         std = (2.0 / (in_features + out_features)) ** 0.5
         self.weights = nn.Parameter(torch.randn(n_angular, out_features, in_features) * std)
 
         self.bias = nn.Parameter(torch.zeros(out_features))
 
-    def forward(self, A_cos, A_sin):
-        A_cos_out = torch.einsum('...id,doi->...od', A_cos, self.weights)
-        A_sin_out = torch.einsum('...id,doi->...od', A_sin, self.weights)
-
-        # Bias only on m=0 (index 0)
-        A_cos_out[..., 0] = A_cos_out[..., 0] + self.bias
+    def _use_dense(self, x):
+        if self.dense_gemm is not None:
+            return self.dense_gemm
+        if not x.is_cuda:
+            return False
+        if x.dtype in (torch.float16, torch.bfloat16):
+            return True
+        return (x.dtype == torch.float32
+                and torch.backends.cuda.matmul.allow_tf32)
+
+    def _dense_weight(self):
+        """(in·n_ang, out·n_ang) block-diagonal-per-m weight, row-major flat
+        index c = feature·n_ang + m on both sides."""
+        na, Fo, Fi = self.weights.shape
+        Wd = self.weights.new_zeros(Fi, na, Fo, na)
+        m = torch.arange(na, device=self.weights.device)
+        Wd[:, m, :, m] = self.weights.permute(0, 2, 1)   # (na, Fi, Fo) slot
+        return Wd.reshape(Fi * na, Fo * na)
 
-        return A_cos_out, A_sin_out
+    def forward(self, A_cos, A_sin):
+        if not self._use_dense(A_cos):
+            A_cos_out = torch.einsum('...id,doi->...od', A_cos, self.weights)
+            A_sin_out = torch.einsum('...id,doi->...od', A_sin, self.weights)
+            A_cos_out[..., 0] = A_cos_out[..., 0] + self.bias   # bias on m=0 only
+            return A_cos_out, A_sin_out
+
+        na, Fo = self.n_angular, self.out_features
+        lead = A_cos.shape[:-2]
+        Wd = self._dense_weight()
+        # Bias enters the GEMM epilogue (addmm), on the m=0 slots of the cos
+        # part only — sin has no invariant channel.
+        bias_full = self.bias.new_zeros(Fo, na)
+        bias_full[:, 0] = self.bias
+        x_cos = A_cos.reshape(-1, self.in_features * na)
+        x_sin = A_sin.reshape(-1, self.in_features * na)
+        y_cos = torch.addmm(bias_full.reshape(-1), x_cos, Wd)
+        y_sin = x_sin @ Wd
+        return y_cos.view(*lead, Fo, na), y_sin.view(*lead, Fo, na)
 
 
 class RealSpaceNonlinearity(nn.Module):

scripts/eval_spice.py

@@ -174,6 +174,10 @@ def main():
     # Eval options
     parser.add_argument('--batch_size',       type=int,   default=8)
     parser.add_argument('--float32',          action='store_true')
+    parser.add_argument('--tf32',             action='store_true',
+                        help='Enable TF32 matmuls (CUDA + --float32 only) — '
+                             'measure the test MAEs at the same precision the '
+                             'model is benchmarked/deployed at.')
     parser.add_argument('--device',           default=None)
     parser.add_argument('--ignore_les',       action='store_true',
                         help='Evaluate the short-range part of an LES checkpoint '
@@ -183,7 +187,15 @@ def main():
     dtype  = torch.float32 if args.float32 else torch.float64
     device = torch.device(args.device if args.device else
                           ('cuda' if torch.cuda.is_available() else 'cpu'))
-    print(f"Device: {device}, dtype: {dtype}")
+    if args.tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+        torch.set_float32_matmul_precision('high')
+        if dtype == torch.float64:
+            print('[tf32] requested but dtype=float64 → no effect '
+                  '(TF32 is float32-only); add --float32 to use it')
+    print(f"Device: {device}, dtype: {dtype}"
+          f"{' +tf32' if args.tf32 and dtype == torch.float32 else ''}")
 
     # ── Load checkpoint ────────────────────────────────────────────────────
     ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False)

scripts/train_ecenet.py

@@ -214,6 +214,7 @@ def train_ecenet(
     eval_batch_size=32,
     seed=42,
     dtype=torch.float64,
+    tf32=False,              # route float32 matmuls to TF32 tensor cores (Ampere+)
     device=None,
     checkpoint_path=None,
     reset_optimizer=False,
@@ -248,6 +249,25 @@ def loss_fn(pred, tgt):
     elif isinstance(device, str):
         device = torch.device(device)
 
+    # TF32: on Ampere+ the step is dominated by fp32 matmuls (forward and the
+    # force double-backward). Routing those to TF32 tensor cores is the cheapest
+    # large speedup, but TF32 keeps only ~10 mantissa bits — A/B the val
+    # force/energy MAE before trusting it. No effect under float64 (TF32 is a
+    # float32-only mode), so warn rather than silently do nothing. (Same
+    # switch as train_ecenet_spice / train_ecenet_xyz.)
+    if tf32:
+        if dtype == torch.float64:
+            if verbose:
+                print_flush("  [tf32] requested but dtype=float64 → no effect "
+                            "(TF32 is float32-only); use dtype=torch.float32")
+        else:
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.allow_tf32 = True
+            torch.set_float32_matmul_precision('high')
+            if verbose:
+                print_flush("  [tf32] enabled: float32 matmuls → TF32 tensor cores "
+                            "(A/B the val MAE against a tf32=False run)")
+
     # ── Data ──────────────────────────────────────────────────────────────
     (pos_train, frc_train, eng_train,
      pos_val,   frc_val,   eng_val,

tests/test_edge_frame_kernel.py

@@ -441,6 +441,24 @@ def test_triton_paths():
     print("test_triton_paths[pack_unrotate]: OK")
 
 
+def test_triton_packed_d():
+    """CUDA-only: run test_triton_paths under BOTH _EF_PACKD settings, so the
+    packed-D kernels (the default) and the table-gather kernels (the
+    ECENET_EF_PACKD=0 fallback) both stay pinned to the fp64 reference."""
+    if not torch.cuda.is_available():
+        print("test_triton_packed_d: SKIP (no CUDA)")
+        return
+    import ecenet.edge_frame_kernel as efk
+    old = efk._EF_PACKD
+    try:
+        for packed in (True, False):
+            efk._EF_PACKD = packed
+            test_triton_paths()
+            print(f"test_triton_packed_d[packed={packed}]: OK")
+    finally:
+        efk._EF_PACKD = old
+
+
 def test_model_integration():
     """Full ECENet (no MP): energy and autograd forces identical flag on/off."""
     import ecenet as _ecenet
@@ -524,6 +542,7 @@ def run():
     test_pack_unrotate_matches_mp_ops()
     test_pack_unrotate_gradchecks()
     test_triton_paths()
+    test_triton_packed_d()
     test_model_integration()
     test_mp_integration()
     print("\nAll tests passed.")

tests/test_equivariant_linear.py

@@ -0,0 +1,158 @@
+"""Tests for the block-diagonal-GEMM EquivariantLinear (``ecenet/equivariant.py``).
+
+The layer computes y[e,o,m] = Σ_i x[e,i,m]·W[m,o,i] (bias on the m=0 cos slot
+only) as one dense row-major GEMM against a per-m block-diagonal weight. These
+tests pin it to the original einsum formulation:
+
+  1. forward equivalence vs the einsum reference (several shapes, fp64);
+  2. gradient equivalence — input, weight, and bias grads vs the reference;
+  3. non-contiguous (m-major, einsum-layout) inputs give the same result;
+  4. batched leading dims (B, n_e, F, n_ang) round-trip unchanged.
+
+Pure PyTorch on CPU (fp64). Run:  python tests/test_equivariant_linear.py
+"""
+
+import os
+import sys  # noqa: E402 — repo root on path for `import ecenet` when run as a script
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+import torch
+
+from ecenet.equivariant import EquivariantLinear
+
+torch.manual_seed(0)
+DTYPE = torch.float64
+
+
+def reference(A_cos, A_sin, weights, bias):
+    """The original einsum formulation — the spec."""
+    oc = torch.einsum('...id,doi->...od', A_cos, weights)
+    os_ = torch.einsum('...id,doi->...od', A_sin, weights)
+    oc = oc.clone()
+    oc[..., 0] = oc[..., 0] + bias
+    return oc, os_
+
+
+def make_layer(Fi=12, Fo=9, m_max=2, seed=0):
+    torch.manual_seed(seed)
+    lin = EquivariantLinear(Fi, Fo, m_max + 1, m_max).to(DTYPE)
+    lin.dense_gemm = True     # force the dense path (auto picks einsum on CPU)
+    with torch.no_grad():                # non-trivial bias
+        lin.bias.add_(torch.randn_like(lin.bias))
+    return lin
+
+
+def make_inputs(n_e=40, Fi=12, m_max=2, seed=1):
+    g = torch.Generator().manual_seed(seed)
+    mk = lambda: torch.randn(n_e, Fi, m_max + 1, generator=g, dtype=DTYPE)  # noqa: E731
+    return mk(), mk()
+
+
+def test_forward_equivalence():
+    worst = 0.0
+    for Fi, Fo, m_max in ((12, 9, 2), (7, 7, 3), (1, 5, 0), (16, 3, 1)):
+        lin = make_layer(Fi, Fo, m_max)
+        A_cos, A_sin = make_inputs(Fi=Fi, m_max=m_max)
+        oc, os_ = lin(A_cos, A_sin)
+        rc, rs = reference(A_cos, A_sin, lin.weights, lin.bias)
+        e = max((oc - rc).abs().max(), (os_ - rs).abs().max()).item()
+        worst = max(worst, e)
+        assert e < 1e-12, f"forward Fi={Fi},Fo={Fo},m={m_max}: {e:.2e}"
+    print(f"  forward equivalence vs einsum reference (worst {worst:.1e})")
+
+
+def test_gradient_equivalence():
+    lin = make_layer()
+    A_cos, A_sin = make_inputs(seed=2)
+    goc = torch.randn_like(A_cos[:, :9])   # (n_e, Fo, n_ang)
+    gos = torch.randn_like(A_cos[:, :9])
+
+    def grads(fn):
+        a = A_cos.detach().clone().requires_grad_(True)
+        b = A_sin.detach().clone().requires_grad_(True)
+        lin.zero_grad()
+        oc, os_ = fn(a, b)
+        (oc * goc + os_ * gos).sum().backward()
+        return a.grad, b.grad, lin.weights.grad.clone(), lin.bias.grad.clone()
+
+    new = grads(lambda a, b: lin(a, b))
+    ref = grads(lambda a, b: reference(a, b, lin.weights, lin.bias))
+    worst = 0.0
+    for name, n, r in zip(("dA_cos", "dA_sin", "dW", "db"), new, ref):
+        e = (n - r).abs().max().item()
+        worst = max(worst, e)
+        assert e < 1e-12, f"{name} mismatch: {e:.2e}"
+    print(f"  gradient equivalence: input/weight/bias grads match (worst {worst:.1e})")
+
+
+def test_noncontiguous_input():
+    """m-major (einsum-layout) inputs — strides (F, 1, n_e·F) — must give the
+    same result as contiguous ones (reshape copies as needed internally)."""
+    lin = make_layer()
+    A_cos, A_sin = make_inputs(seed=3)
+
+    def as_m_major(t):
+        n_e, F, na = t.shape
+        out = torch.empty(na, n_e, F, dtype=t.dtype).permute(1, 2, 0)
+        out.copy_(t)
+        assert not out.is_contiguous()
+        return out
+
+    oc, os_ = lin(A_cos, A_sin)
+    mc, ms = lin(as_m_major(A_cos), as_m_major(A_sin))
+    e = max((oc - mc).abs().max(), (os_ - ms).abs().max()).item()
+    assert e < 1e-15, f"m-major input mismatch: {e:.2e}"
+    print(f"  non-contiguous (m-major) input matches contiguous ({e:.1e})")
+
+
+def test_batched_leading_dims():
+    lin = make_layer()
+    A_cos, A_sin = make_inputs(n_e=24, seed=4)
+    B_cos = A_cos.reshape(4, 6, 12, 3)
+    B_sin = A_sin.reshape(4, 6, 12, 3)
+    oc, os_ = lin(A_cos, A_sin)
+    bc, bs = lin(B_cos, B_sin)
+    assert bc.shape == (4, 6, 9, 3) and bs.shape == (4, 6, 9, 3)
+    e = max((bc.reshape(24, 9, 3) - oc).abs().max(),
+            (bs.reshape(24, 9, 3) - os_).abs().max()).item()
+    assert e < 1e-15, f"batched mismatch: {e:.2e}"
+    print(f"  batched leading dims match flat ({e:.1e})")
+
+
+def test_dispatch():
+    """Auto-dispatch: einsum on CPU / strict fp32 / fp64; dense only where the
+    padded GEMM rides tensor cores (CUDA fp16/bf16, or fp32 with TF32 on).
+    dense_gemm=True/False overrides both ways."""
+    lin = EquivariantLinear(4, 4, 3, 2)
+    x64 = torch.randn(2, 4, 3, dtype=torch.float64)
+    assert not lin._use_dense(x64), "CPU/fp64 must take the einsum path"
+    lin.dense_gemm = True
+    assert lin._use_dense(x64), "dense_gemm=True must override"
+    lin.dense_gemm = False
+    assert not lin._use_dense(x64), "dense_gemm=False must override"
+    lin.dense_gemm = None
+    if torch.cuda.is_available():
+        x32 = torch.randn(2, 4, 3, dtype=torch.float32, device='cuda')
+        old = torch.backends.cuda.matmul.allow_tf32
+        try:
+            torch.backends.cuda.matmul.allow_tf32 = False
+            assert not lin._use_dense(x32), "CUDA fp32 without TF32 → einsum"
+            torch.backends.cuda.matmul.allow_tf32 = True
+            assert lin._use_dense(x32), "CUDA fp32 with TF32 → dense"
+        finally:
+            torch.backends.cuda.matmul.allow_tf32 = old
+        print("  dispatch: einsum on CPU/strict-fp32, dense on TF32; overrides work")
+    else:
+        print("  dispatch: einsum on CPU, overrides work (CUDA cases skipped)")
+
+
+if __name__ == "__main__":
+    print("EquivariantLinear block-diagonal-GEMM tests")
+    test_forward_equivalence()
+    test_gradient_equivalence()
+    test_noncontiguous_input()
+    test_batched_leading_dims()
+    test_dispatch()
+    print("All tests passed.")

tools/benchmark_calculator.py

@@ -169,6 +169,13 @@ def main():
                    help='replicate the box into an AxBxC supercell (size scaling)')
     p.add_argument('--float32', action='store_true',
                    help='float32 where the factory supports it (state it in the table)')
+    p.add_argument('--tf32', action='store_true',
+                   help='enable TF32 matmuls process-wide (torch backends flag; '
+                        'float32 only). ecenet: also required for the dense '
+                        'EquivariantLinear GEMM path to auto-dispatch. Recorded '
+                        'in the CSV dtype column. Note --calc dpa ignores torch '
+                        "flags — DeePMD's precision is set via DP_TF32_INFER "
+                        'and friends in its own environment.')
     p.add_argument('--fuse', action='store_true',
                    help="enable the model's fast paths (ecenet: edge-frame + "
                         'activation fusion; mace: enable_cueq; no-op for '
@@ -186,6 +193,16 @@ def main():
     p.add_argument('--csv', default=None, help='append one result row (collate runs)')
     args = p.parse_args()
 
+    if args.tf32:
+        if torch is None:
+            sys.exit('--tf32 needs torch installed')
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+        torch.set_float32_matmul_precision('high')
+        if not args.float32:
+            print('[tf32] set, but without --float32 most factories run float64 '
+                  '— TF32 will not apply there')
+
     atoms = read(args.box, index=args.frame_idx)
     if args.repeat:
         atoms = atoms.repeat(tuple(args.repeat))
@@ -248,7 +265,8 @@ def main():
                'checkpoint': args.checkpoint or '', 'box': os.path.basename(args.box),
                'repeat': 'x'.join(map(str, args.repeat)) if args.repeat else '',
                'n_atoms': n_atoms,
-               'dtype': 'float32' if args.float32 else 'default/float64',
+               'dtype': (('float32' if args.float32 else 'default/float64')
+                         + ('+tf32' if args.tf32 else '')),
                'fuse': args.fuse, 'device': gpu,
                'host': os.uname().nodename,  # login vs compute node — shared-GPU
                                              # login timings are contended garbage