Tutorial 9: Accelerating the Recurrent Matvec

Note

Reading Time: 20-25 minutes

Difficulty: Intermediate

Prerequisites: Tutorial 1 (Building and Using CANN Models), familiarity with NumPy/JAX matmuls

This tutorial covers the strategy pattern that powers CANN1D / CANN2D acceleration: the model holds a single strategy object on self.irec_backend, and update() just calls it. By the end you will know which of the five accl_mode values to pick for your workload, when the FFT path applies, and how to add a new backend of your own.

---

1. Why accelerate?

The CANN dynamics spend almost all of their time in one line:

Irec = self._accel_Irec(self.r.value)  # → self.irec_backend(r)

In dense mode that is a full W @ r matvec, which is O(n²) in the number of neurons. For CANN1D at num = 4096 that is a 16M-element multiplication every step, and it dominates everything else. The five accl_mode values cover different ways to make this cheaper — either by approximation (low-rank SVD) or by exploiting hidden structure (circulant FFT).

# Quick demonstration: time the dense matvec at increasing num.
import time
import brainpy.math as bm
import numpy as np
from canns.models.basic import CANN1D

bm.set_dt(0.1)
results = []
for num in [128, 512, 2048, 4096]:
    m = CANN1D(num=num, accl_mode="normal")
    r = bm.random.rand(num)
    # Warm up
    for _ in range(3):
        m.irec_backend(r)
    t0 = time.perf_counter()
    for _ in range(30):
        m.irec_backend(r)
    results.append((num, (time.perf_counter() - t0) / 30 * 1e6))

print(f"{'num':>6}  {'dense (μs/step)':>16}")
for num, t in results:
    print(f"{num:>6}  {t:>16.1f}")

   num   dense (μs/step)
   128               7.6
   512               6.7
  2048              10.7
  4096              37.0

The dense matvec grows roughly like num²: 128 → 512 is a 4× wall-clock jump, and 512 → 2048 is another 16×. Anything we can do to bring this down to O(n · k) or O(n log n) will pay off fast.

2. The strategy pattern: self.irec_backend

Every accelerated model carries one backend object:

model.irec_backend        # ← the strategy
model.irec_backend(r)     # ← the call site (this is what update() does)

The model class does not know which mode is active. The backend is responsible for both the pre-compute (SVD, FFT) and the per-step matvec. Adding a new acceleration mode is a pure addition to canns.models.basic.accel — cann.py does not need to change.

from canns.models.basic import CANN1D
import brainpy.math as bm

bm.set_dt(0.1)
m = CANN1D(num=128, accl_mode="fast")
print(f"  accl_mode:    {m.accl_mode}")
print(f"  accl_k:       {m.accl_k}")
print(f"  is_accelerated: {m.is_accelerated}")
print(f"  backend type: {type(m.irec_backend).__name__}")
print(f"  backend.mode: {m.irec_backend.mode}")
print(f"  backend.k:    {m.irec_backend.k}")
print(f"  U shape:      {m.irec_backend.U.shape}")
print(f"  V shape:      {m.irec_backend.V.shape}")

  accl_mode:    fast
  accl_k:       8
  is_accelerated: True
  backend type: LowRankIrec
  backend.mode: fast
  backend.k:    8
  U shape:      (128, 8)
  V shape:      (128, 8)

The public accl_mode / accl_k / is_accelerated properties are thin delegates to self.irec_backend; the test suite also reads m._U_l and m._V_l, which are properties that look up self.irec_backend.U / .V for inspection.

3. Decision table — which mode should I use?

The same content as the table at the top of canns.models.basic.accel, but rendered as a notebook-friendly list-table.

accl_mode	Math	Exact?	n=2048 speedup	When to use
"normal"	W @ r	Yes	1×	Small networks (num <= 128), or whenever you need the textbook reference.
"fast"	U @ (Vᵀ @ r), k from defaults	Approx (sub-mrad)	30–245×	Recommended for production. Default ranks keep bump-position error below ~5 mrad.
"ultra-fast"	U @ (Vᵀ @ r) at k=1 (1D) or k=4 (2D)	Coarser	100–500×	When the bump only needs to exist, not be positioned precisely.
"auto"	rank picked from SVD spectrum to satisfy accl_target_err_mrad	Approx (≤ budget)	varies	When you want to set the error budget (mrad) instead of the rank.
"fft"	real(ifft(K_fft ⊙ fft(r))) (1D) / 2D analogue	Yes	25–50×	Only on a clean ring/torus with endpoint=False. Falls back to dense with a warning otherwise.

4. Mode walk-through

4.1 "normal" — full-rank reference

No pre-compute. W @ r every step. The reference for accuracy comparisons, and the fastest option for small networks where the constant-factor overhead of any approximation is not worth it.

m = CANN1D(num=128, accl_mode="normal")
print(f"  type: {type(m.irec_backend).__name__}")
print(f"  k={m.accl_k}  (sentinel for full-rank / no approximation)")
print(f"  U is None: {m.irec_backend.U is None}")

  type: DenseIrec
  k=-1  (sentinel for full-rank / no approximation)
  U is None: True

4.2 "fast" — recommended for production

Truncated SVD at the default rank (CANN1D k=8, CANN2D k=32). The Gaussian distance kernel that make_conn builds is well-approximated by a low-rank matrix, so 8 components capture 99%+ of the singular-value energy. See benchmarks/canns-accl/lowrank for the full sweep.

m = CANN1D(num=2048, accl_mode="fast")
print(f"  type: {type(m.irec_backend).__name__}")
print(f"  k={m.accl_k}  (default for CANN1D fast)")
print(f"  factor shapes: U={m.irec_backend.U.shape}, V={m.irec_backend.V.shape}")

  type: LowRankIrec
  k=8  (default for CANN1D fast)
  factor shapes: U=(2048, 8), V=(2048, 8)

4.3 "ultra-fast" — k=1 / k=4 only

For when the bump only needs to exist. The leading singular vector of a Gaussian kernel is itself a Gaussian, so k=1 already preserves bump tracking (just not bump width). At k=1 the bump is wider and the position error is a few mrad — useful for very large num / length or for headroom in production.

m = CANN1D(num=2048, accl_mode="ultra-fast")
print(f"  type: {type(m.irec_backend).__name__}")
print(f"  k={m.accl_k}  (ultra-fast = k=1 for CANN1D)")

  type: LowRankIrec
  k=1  (ultra-fast = k=1 for CANN1D)

4.4 "auto" — error-budget-driven rank

You set the error budget (accl_target_err_mrad) and the factory picks the smallest rank whose predicted bump-tracking error fits. The heuristic is calibrated against the canns_lowrank benchmark and is documented in canns.models.basic.cann._pick_k_for_err_target().

m = CANN1D(num=2048, accl_mode="auto", accl_target_err_mrad=0.5)
print(f"  mode:    {m.accl_mode}")
print(f"  picked k={m.accl_k}  (for a 0.5 mrad budget)")

# Tightening the budget at runtime picks a higher k.
m.set_accl_mode("auto", target_err_mrad=0.1)
print(f"  after tightening to 0.1 mrad: k={m.accl_k}")

  mode:    auto
  picked k=5  (for a 0.5 mrad budget)
  after tightening to 0.1 mrad: k=8

4.5 "fft" — exact circulant matvec

The Gaussian distance kernel on a uniform ring is right-circulant, which means its matvec is a circular convolution — exactly what the convolution theorem says can be computed in O(n log n) via FFT. This is not an approximation; on a clean circulant the result equals W @ r to fp32 precision.

The catch: the canns default endpoint=True grid is not circulant (the wrap point duplicates the endpoint), so requesting accl_mode="fft" on a default-built model falls back to dense with a warning. The fix is to override the grid to endpoint=False and rebuild conn_mat before switching to FFT mode.

# 1) Default grid (endpoint=True) — falls back to dense with a warning.
import warnings
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    m = CANN1D(num=128, accl_mode="fft")
print(f"  requested: fft → got: {m.accl_mode}  (silent fallback)")
print(f"  warning: {[str(w.message) for w in caught][0][:80]}...")

  requested: fft → got: normal  (silent fallback)
  warning: accl_mode='fft' requested but the grid appears to use endpoint=True (x[0] = x[-1...

# 2) Clean grid (endpoint=False) — the FFT path becomes exact.
m = CANN1D(num=128, accl_mode="normal")
m.x = bm.linspace(-bm.pi, bm.pi, 128, endpoint=False)
m.conn_mat = m.make_conn()
m.set_accl_mode("fft")
print(f"  mode: {m.accl_mode}, K_fft shape: {m.irec_backend.K_fft.shape}")
print(f"  backend type: {type(m.irec_backend).__name__}")

  mode: fft, K_fft shape: (128,)
  backend type: CirculantFFTIrec1D

# Verify: the FFT matvec equals the dense matvec to ~1e-5 on the clean grid.
r = bm.random.rand(128)
dense = np.asarray(m.conn_mat @ r)
fft_out = np.asarray(m.irec_backend(r))
print(f"  max |dense - fft|: {np.abs(dense - fft_out).max():.2e}")

  max |dense - fft|: 1.14e-05

5. End-to-end dynamics: low-rank ≈ dense

The whole point of approximation is to be approximately right. We confirm that a "fast" model produces the same r_max trajectory as a "normal" model under a moving stimulus.

bm.set_dt(0.1)
num = 256
T = 80
m_n = CANN1D(num=num, accl_mode="normal")
m_f = CANN1D(num=num, accl_mode="fast")
x = np.asarray(m_n.x)
z_range = float(m_n.z_range)

def run(m):
    traj = np.empty((T, num), dtype=np.float32)
    for t in range(T):
        pos = np.pi * t / (T - 1)
        d = (x - pos) % z_range
        d = np.where(d > z_range / 2, d - z_range, d)
        inp = (m.A * np.exp(-0.25 * (d / m.a) ** 2)).astype(np.float32)
        m.update(inp)
        traj[t] = np.asarray(m.r.value)
    return traj

r_n, r_f = run(m_n), run(m_f)
print(f"  r_max trajectory max diff: {np.abs(r_n.max(1) - r_f.max(1)).max():.2e}")
print(f"  ✓ (the fast and dense models agree to ~1e-7 over 80 steps)")

  r_max trajectory max diff: 9.53e-07
  ✓ (the fast and dense models agree to ~1e-7 over 80 steps)

6. Runtime mode switching

Because the backend is just an attribute, you can swap strategies at runtime via set_accl_mode. This is useful for “start dense, then accelerate” workflows (e.g. settle a bump with the reference dynamics, then switch to FFT to run a long sequence).

bm.set_dt(0.1)
m = CANN1D(num=256, accl_mode="normal")
print(f"  start:       mode={m.accl_mode}  backend={type(m.irec_backend).__name__}")
m.set_accl_mode("fast")
print(f"  → fast:      mode={m.accl_mode}  backend={type(m.irec_backend).__name__}  k={m.accl_k}")
m.set_accl_mode("auto", target_err_mrad=0.5)
print(f"  → auto:      mode={m.accl_mode}  backend={type(m.irec_backend).__name__}  k={m.accl_k}")
m.set_accl_mode("normal")
print(f"  → normal:    mode={m.accl_mode}  backend={type(m.irec_backend).__name__}")

  start:       mode=normal  backend=DenseIrec
  → fast:      mode=fast  backend=LowRankIrec  k=8
  → auto:      mode=auto  backend=LowRankIrec  k=5
  → normal:    mode=normal  backend=DenseIrec

7. Adding a new backend

Three steps, all in canns.models.basic.accel. cann.py is not touched.

1. Write a class that carries the pre-computed factors you need and implements __call__(self, r):

class BlockCirculantIrec:
    mode = "block-circulant"
    k = -1

    def __init__(self, blocks):
        self.blocks = blocks  # list of (W, U_fft, V_fft) per block

    def __call__(self, r):
        # your math here
        ...

2. Register the mode in ACCL_MODES (and in ACCL_DEFAULT_K if it has per-family defaults):

ACCL_MODES = ("normal", "fast", "ultra-fast", "auto", "fft", "block-circulant")

3. Add a dispatch branch in make_irec_backend():

if mode == "block-circulant":
    return BlockCirculantIrec(_compute_block_factors(model.conn_mat))

The model sees a new callable on self.irec_backend; everything downstream (update(), the property delegates, set_accl_mode) keeps working.

8. References

• benchmarks/canns-accl/lowrank/ — accuracy / speedup sweep across the four low-rank modes at multiple num and length values. The default ranks in ACCL_DEFAULT_K come from this benchmark.

• benchmarks/canns-accl/fft/ — exact vs approximate comparison for the FFT mode on clean and default grids.

• canns.models.basic.cann._pick_k_for_err_target() — the spectrum-based rank picker used by "auto" mode, with its calibration table.

• canns.models.basic.accel — the strategy module itself; the same decision table at the top of the module is the canonical reference for “which mode should I use?”.