教程 9: 加速循环矩阵-向量乘¶
Note
阅读时间：20-25 分钟
难度：中级
先修：教程 1（CANN 模型的构建与使用），熟悉 NumPy / JAX 矩阵乘法
本教程介绍支撑 CANN1D / CANN2D 加速的「策略模式」：模型在 self.irec_backend 上持有一个策略对象，update() 只需要调用它。读完本节你将知道五种 accl_mode 如何为你的工作负载选择、FFT 路径何时适用，以及如何自己添加一个后端。
1. 为什么要加速？¶
CANN 动力学几乎所有时间都花在这一行上：:
Irec = self._accel_Irec(self.r.value) # → self.irec_backend(r)
在 dense 模式下，这是一个完整的 W @ r 矩阵-向量乘，复杂度为 O(n²)。对于 num = 4096 的 CANN1D，每一步都是 16M 元素乘法，而且它会主宰所有其他开销。五个 accl_mode 值提供了让这一步更便宜的不同方式——要么通过近似（低秩 SVD），要么利用隐藏结构（循环 FFT）。
[1]:
# 快速演示：在不同的 num 下计时 dense 矩阵-向量乘。
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 9.7
2048 8.4
4096 30.9
dense 矩阵-向量乘大致以 num² 增长：128 → 512 是 4× 的耗时跳变，512 → 2048 又跳 16×。任何能把它降到 O(n · k) 或 O(n log n) 的方法都会很快见效。
2. 策略模式：self.irec_backend¶
每个被加速的模型持有一个后端对象：:
model.irec_backend # ← 策略
model.irec_backend(r) # ← 调用点（也就是 update() 干的事）
模型类并不「知道」当前是哪个模式。后端既负责预计算（SVD、FFT），也负责每一步的矩阵-向量乘。新增一个加速模式是一种**纯加法**——只在 canns.models.basic.accel 里加，cann.py 不需要任何改动。
[2]:
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)
公开的 accl_mode / accl_k / is_accelerated 属性是 self.irec_backend 的薄代理；测试套件也读 m._U_l 和 m._V_l，这些属性把请求转发到 self.irec_backend.U / .V 用于内省。
3. 决策表——我该用哪个模式？¶
内容跟 canns.models.basic.accel 顶部的决策表一样，但用 list-table 渲染更友好。
|
数学形式 |
精确？ |
n=2048 加速比 |
何时使用 |
|---|---|---|---|---|
|
|
精确 |
1× |
小网络（ |
|
|
近似（亚 mrad） |
30–245× |
生产推荐。默认 k 下 bump 位置误差 < 5 mrad。 |
|
|
较粗 |
100–500× |
只需 bump 存在、对位置精度不敏感时。 |
|
根据 |
近似（不超过预算） |
不定 |
想设置「误差预算」(mrad) 而非具体 rank。 |
|
``real(ifft(K_fft ⊙ fft(r)))``（1D）/ 2D 类比 |
精确 |
25–50× |
**只能**用于 |
4. 各模式详解¶
4.1 "normal" —— 全秩参考实现¶
无预计算。每一步 W @ r。是精度对比的参考；当网络小到任何近似都不划算时也是最快的选项。
[3]:
m = CANN1D(num=128, accl_mode="normal")
print(f" type: {type(m.irec_backend).__name__}")
print(f" k={m.accl_k} (full-rank / 无近似的哨兵值)")
print(f" U is None: {m.irec_backend.U is None}")
type: DenseIrec
k=-1 (full-rank / 无近似的哨兵值)
U is None: True
4.2 "fast" —— 生产推荐¶
在默认 rank（CANN1D k=8，CANN2D k=32）处截断 SVD。make_conn 构造的高斯距离核可以用低秩矩阵很好地逼近，8 个分量就足以捕获 99% 以上的奇异值能量。完整扫描见 benchmarks/canns-accl/lowrank。
[4]:
m = CANN1D(num=2048, accl_mode="fast")
print(f" type: {type(m.irec_backend).__name__}")
print(f" k={m.accl_k} (CANN1D fast 的默认)")
print(f" factor shapes: U={m.irec_backend.U.shape}, V={m.irec_backend.V.shape}")
type: LowRankIrec
k=8 (CANN1D fast 的默认)
factor shapes: U=(2048, 8), V=(2048, 8)
4.3 "ultra-fast" —— 仅 k=1 / k=4¶
用于只需 bump 存在的场景。高斯核的主导奇异向量本身就是一个高斯，因此 k=1 就足以保留 bump 的**追踪**（只是 bump 宽度会变宽）。在 k=1 下 bump 较宽、位置误差是几个 mrad——非常 num / length 很大的情况下有用，或生产中需要更宽裕度时。
[5]:
m = CANN1D(num=2048, accl_mode="ultra-fast")
print(f" type: {type(m.irec_backend).__name__}")
print(f" k={m.accl_k} (CANN1D ultra-fast = k=1)")
type: LowRankIrec
k=1 (CANN1D ultra-fast = k=1)
4.4 "auto" —— 误差预算驱动的 rank 选择¶
你设置**误差预算**（accl_target_err_mrad），工厂会选最小 rank 使得预测的 bump 追踪误差不超过预算。启发式与 canns_lowrank 基准对齐，详见 canns.models.basic.cann._pick_k_for_err_target()。
[6]:
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} (对应 0.5 mrad 预算)")
# 运行时收紧预算会选更高 k。
m.set_accl_mode("auto", target_err_mrad=0.1)
print(f" 收紧到 0.1 mrad 后: k={m.accl_k}")
mode: auto
picked k=5 (对应 0.5 mrad 预算)
收紧到 0.1 mrad 后: k=8
4.5 "fft" —— 精确的循环矩阵-向量乘¶
在均匀环上的高斯距离核是右循环的，意味着它的矩阵-向量乘是循环卷积——这正是卷积定理说可以由 FFT 在 O(n log n) 内完成的运算。这并非近似：在干净循环阵上，结果在 fp32 精度内等于 W @ r。
注意：canns 默认的 endpoint=True 网格并不是循环的（环接点会复制端点），因此在默认模型上请求 accl_mode="fft" 会回退到 dense 并告警。修复方法是先把网格覆盖成 endpoint=False 并重建 conn_mat，再切换到 FFT 模式。
[7]:
# 1) 默认网格（endpoint=True）—— 回退到 dense 并告警。
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} (静默回退)")
print(f" warning: {[str(w.message) for w in caught][0][:80]}...")
requested: fft → got: normal (静默回退)
warning: accl_mode='fft' requested but the grid appears to use endpoint=True (x[0] = x[-1...
[8]:
# 2) 干净网格（endpoint=False）—— FFT 路径变为精确。
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
[9]:
# 验证：在干净网格上，FFT 矩阵-向量乘与 dense 在 ~1e-5 量级一致。
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.53e-05
5. 端到端动力学：低秩 ≈ dense¶
近值的全部意义就是「近似正确」。下面我们验证在移动刺激下，"fast" 模型与 "normal" 模型产生相同的 r_max 轨迹。
[10]:
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" ✓ (fast 与 dense 在 80 步内一致到 ~1e-7)")
r_max trajectory max diff: 9.53e-07
✓ (fast 与 dense 在 80 步内一致到 ~1e-7)
6. 运行时切换模式¶
因为后端只是一个属性，所以你可以通过 set_accl_mode 在运行时替换策略。这对「先 dense 启动、然后加速」的工作流很有用（例如先用参考动力学稳定一个 bump，再切到 FFT 跑一段长序列）。
[11]:
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. 添加新后端¶
三步搞定，都改在 canns.models.basic.accel 里。cann.py 不用动。
写一个类，持有所需的预计算因子并实现
__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):
# 在这里写你的数学
...
**在
ACCL_MODES中注册该模式**（如果它有 per-family 默认 rank，也加到ACCL_DEFAULT_K）：
ACCL_MODES = ("normal", "fast", "ultra-fast", "auto", "fft", "block-circulant")
在 :func:`make_irec_backend` 里加一个分发分支：
if mode == "block-circulant":
return BlockCirculantIrec(_compute_block_factors(model.conn_mat))
模型在 self.irec_backend 上看到一个新 callable；下游所有逻辑（update()、属性代理、set_accl_mode）都照旧工作。
8. 参考资料¶
benchmarks/canns-accl/lowrank/—— 跨四个低秩模式、多个num和length的精度 / 加速扫描。ACCL_DEFAULT_K中的默认 rank 来自此基准。benchmarks/canns-accl/fft/—— FFT 模式在干净 / 默认网格下的精确 vs 近似对比。canns.models.basic.cann._pick_k_for_err_target()——"auto"模式所用的基于谱的 rank 选择器，含校准表。canns.models.basic.accel—— 策略模块本身；模块顶部那张决策表是「我该用哪个模式？」的官方参考。