925 lines
35 KiB
Plaintext
925 lines
35 KiB
Plaintext
# -*- coding: utf-8 -*-
|
|
# ============================================================
|
|
# torch-lite: 仿 PyTorch 的极简张量/神经网络库（Viper）
|
|
#
|
|
# 设计原则:
|
|
# * "足够用就行": 只覆盖 MNIST 两层 MLP 训练/推理所需的原语。
|
|
# * API 命名贴近 PyTorch: Tensor / linear_fwd(bias 融合) / relu /
|
|
# relu_backward / softmax_crossentropy(前向+梯度一体) / sgd_step。
|
|
# * 张量为 2D 稠密 double (rows x cols, 行主序), 权重形状同 PyTorch:
|
|
# Linear(in, out) -> weight (out, in) / bias (out,), forward = x @ W^T + b。
|
|
# * 计算热点结合 vipersimd:
|
|
# - linear_fwd 内积按 k%4 用 simd_dot4d 累加(纯 AVX2 点积);
|
|
# - relu_fwd 用 simd_max4d(x, 0) 四路取大;
|
|
# - 向量乘标量用 simd_mul_scalar4d(归一化 /255 等)。
|
|
# backward / softmax / sgd 等按标量写(数据量小, 可读性优先)。
|
|
# * 内存: 所有 Tensor 及 data/grad 都从调用方传入的 memhub.MemBuddy
|
|
# 分配(与 includes 既有约定一致, pool.alloc 免手动 free)。
|
|
# ============================================================
|
|
from stdint import *
|
|
import t, c
|
|
import stdlib
|
|
import string
|
|
import stdio
|
|
import math
|
|
import vipersimd
|
|
import memhub
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# Tensor
|
|
# ------------------------------------------------------------
|
|
novirtual class Tensor:
|
|
Data: t.CPtr # rows*cols*8 字节连续 double
|
|
Grad: t.CPtr # 同尺寸梯度(可为 None)
|
|
Rows: sizet
|
|
Cols: sizet
|
|
|
|
# 字段: Data/Grad/Rows/Cols = 32 字节(无 vtable)
|
|
TENSOR_SIZE: t.CDefine = 32
|
|
|
|
|
|
def _zeros_block(pool: memhub.MemBuddy*, bytes_n: sizet) -> t.CPtr:
|
|
"""分配并清零一块内存"""
|
|
p: t.CPtr = pool.alloc(bytes_n)
|
|
if p is not None:
|
|
string.memset(p, 0, bytes_n)
|
|
return p
|
|
|
|
|
|
def zeros(pool: memhub.MemBuddy*, rows: sizet,
|
|
cols: sizet, with_grad: int = 1) -> Tensor | t.CPtr:
|
|
"""torch.zeros: 创建 rows x cols 的全零 Tensor, 可选分配 Grad 缓冲"""
|
|
if pool is None:
|
|
return None
|
|
t0: rawptr[Tensor] = pool.alloc(TENSOR_SIZE)
|
|
if t0 is None:
|
|
return None
|
|
string.memset(t0, 0, TENSOR_SIZE)
|
|
nb: sizet = rows * cols * sizet(8)
|
|
t0.Data = _zeros_block(pool, nb)
|
|
if t0.Data is None:
|
|
return None
|
|
if with_grad != 0:
|
|
t0.Grad = _zeros_block(pool, nb)
|
|
t0.Rows = rows
|
|
t0.Cols = cols
|
|
return t0
|
|
|
|
|
|
def from_raw(pool: memhub.MemBuddy*, data: t.CPtr,
|
|
rows: sizet, cols: sizet,
|
|
with_grad: int = 0) -> Tensor | t.CPtr:
|
|
"""torch.from_raw: 包装一块已有内存为 Tensor(不复制, 不拥有)"""
|
|
if pool is None:
|
|
return None
|
|
t0: rawptr[Tensor] = pool.alloc(TENSOR_SIZE)
|
|
if t0 is None:
|
|
return None
|
|
string.memset(t0, 0, TENSOR_SIZE)
|
|
t0.Data = data
|
|
t0.Grad = None
|
|
t0.Rows = rows
|
|
t0.Cols = cols
|
|
return t0
|
|
|
|
|
|
def tensor_free(t0: Tensor | t.CPtr):
|
|
"""释放 Tensor 外壳(数据由 MemBuddy 池统一回收, 无需手动 free)"""
|
|
pass
|
|
|
|
|
|
def tensor_data(t0: Tensor | t.CPtr) -> t.CPtr:
|
|
if t0 is None:
|
|
return None
|
|
return t0.Data
|
|
|
|
|
|
def tensor_grad(t0: Tensor | t.CPtr) -> t.CPtr:
|
|
if t0 is None:
|
|
return None
|
|
return t0.Grad
|
|
|
|
|
|
def tensor_rows(t0: Tensor | t.CPtr) -> sizet:
|
|
if t0 is None:
|
|
return 0
|
|
return t0.Rows
|
|
|
|
|
|
def tensor_cols(t0: Tensor | t.CPtr) -> sizet:
|
|
if t0 is None:
|
|
return 0
|
|
return t0.Cols
|
|
|
|
|
|
def fill(t0: Tensor | t.CPtr, val: double):
|
|
"""torch.fill: 全部元素置为 val"""
|
|
if t0 is None or t0.Data is None:
|
|
return
|
|
d: double | t.CPtr = t0.Data
|
|
n: sizet = t0.Rows * t0.Cols
|
|
i: sizet = 0
|
|
while i < n:
|
|
d[i] = val
|
|
i += 1
|
|
|
|
|
|
def zero_grad(t0: Tensor | t.CPtr):
|
|
"""torch.zero_grad: Grad 缓冲清零"""
|
|
if t0 is None or t0.Grad is None:
|
|
return
|
|
string.memset(t0.Grad, 0, t0.Rows * t0.Cols * sizet(8))
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# 确定性伪随机(LCG, minstd)与 Kaiming 均匀初始化
|
|
# ------------------------------------------------------------
|
|
_rng_state: double = 12345.0
|
|
LCG_M: double = 2147483647.0
|
|
LCG_A: double = 48271.0
|
|
|
|
|
|
def manual_seed(seed: int64):
|
|
global _rng_state
|
|
_rng_state = double(seed % int64(2147483647))
|
|
if _rng_state <= double(0.0):
|
|
_rng_state = double(12345.0)
|
|
|
|
|
|
def _next_rand() -> double:
|
|
"""(0,1) 均匀随机"""
|
|
global _rng_state
|
|
_rng_state = math.fmod(_rng_state * LCG_A, LCG_M)
|
|
return _rng_state / LCG_M
|
|
|
|
|
|
def uniform_fill(t0: Tensor | t.CPtr, lo: double, hi: double):
|
|
"""U(lo, hi) 填充"""
|
|
if t0 is None or t0.Data is None:
|
|
return
|
|
d: double | t.CPtr = t0.Data
|
|
n: sizet = t0.Rows * t0.Cols
|
|
span: double = hi - lo
|
|
i: sizet = 0
|
|
while i < n:
|
|
d[i] = lo + _next_rand() * span
|
|
i += 1
|
|
|
|
|
|
def kaiming_uniform_fill(w: Tensor | t.CPtr, fan_in: sizet,
|
|
fan_out: sizet):
|
|
"""PyTorch Linear 默认初始化: U(-sqrt(6/(fan_in+fan_out)), +...)"""
|
|
bound: double = math.sqrt(
|
|
double(6.0) / double(fan_in + fan_out))
|
|
uniform_fill(w, -bound, bound)
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# 基础元素级运算(结合 vipersimd 4 宽指令)
|
|
# ------------------------------------------------------------
|
|
def vector_mul_scalar4(src: t.CPtr, scalar_ptr: t.CPtr, dst: t.CPtr,
|
|
n: sizet):
|
|
"""dst = src * scalar (逐 4 元素 simd_mul_scalar4d + 标量尾数)"""
|
|
if src is None or scalar_ptr is None or dst is None:
|
|
return
|
|
i: sizet = 0
|
|
while i + 4 <= n:
|
|
vipersimd.simd_mul_scalar4d(
|
|
t.CPtr(uint64(src) + uint64(i) * 8),
|
|
scalar_ptr,
|
|
t.CPtr(uint64(dst) + uint64(i) * 8))
|
|
i += 4
|
|
sd: double | t.CPtr = src
|
|
dd: double | t.CPtr = dst
|
|
# 尾部标量
|
|
sp: double | t.CPtr = scalar_ptr
|
|
sval: double = sp[0]
|
|
while i < n:
|
|
dd[i] = sd[i] * sval
|
|
i += 1
|
|
|
|
|
|
def vector_add_inplace(a: t.CPtr, b: t.CPtr, n: sizet):
|
|
"""a += b (逐 4 元素 simd_add_array 尾部标量)"""
|
|
if a is None or b is None:
|
|
return
|
|
tmp: t.CPtr = stdlib.malloc(8)
|
|
if tmp is None:
|
|
return
|
|
i: sizet = 0
|
|
while i + 4 <= n:
|
|
vipersimd.simd_add_array(
|
|
t.CPtr(uint64(a) + uint64(i) * 8),
|
|
t.CPtr(uint64(b) + uint64(i) * 8),
|
|
t.CPtr(uint64(a) + uint64(i) * 8))
|
|
i += 4
|
|
da: double | t.CPtr = a
|
|
db: double | t.CPtr = b
|
|
while i < n:
|
|
da[i] = da[i] + db[i]
|
|
i += 1
|
|
stdlib.free(tmp)
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# 线性层: forward = x @ W^T + b (W: out x in, x: m x in, out: m x out)
|
|
# 内积按 k%4 用 vipersimd.simd_dot4d —— MNIST 主计算热点
|
|
# ------------------------------------------------------------
|
|
def linear_fwd(w: Tensor | t.CPtr, b: Tensor | t.CPtr,
|
|
x: Tensor | t.CPtr, out: Tensor | t.CPtr):
|
|
if w is None or b is None or x is None or out is None:
|
|
return
|
|
m: sizet = x.Rows
|
|
k: sizet = x.Cols
|
|
n: sizet = w.Rows
|
|
wd: double | t.CPtr = w.Data
|
|
bd: double | t.CPtr = b.Data
|
|
xd: double | t.CPtr = x.Data
|
|
od: double | t.CPtr = out.Data
|
|
acc_buf: t.CPtr = stdlib.malloc(8)
|
|
if acc_buf is None:
|
|
return
|
|
accp: double | t.CPtr = acc_buf
|
|
i: sizet = 0
|
|
while i < m:
|
|
j: sizet = 0
|
|
while j < n:
|
|
acc: double = bd[j]
|
|
kk: sizet = 0
|
|
while kk + 4 <= k:
|
|
vipersimd.simd_dot4d(
|
|
t.CPtr(uint64(xd) + uint64(i * k + kk) * 8),
|
|
t.CPtr(uint64(wd) + uint64(j * k + kk) * 8),
|
|
acc_buf)
|
|
acc = acc + accp[0]
|
|
kk += 4
|
|
while kk < k:
|
|
acc = acc + xd[i * k + kk] * wd[j * k + kk]
|
|
kk += 1
|
|
od[i * n + j] = acc
|
|
j += 1
|
|
i += 1
|
|
stdlib.free(acc_buf)
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# 线性层反向: 输入 w(out,in) x(m,in) g(m,out)
|
|
# dw[j,i] += sum_b g[b,j]*x[b,i] ; db[j] += sum_b g[b,j]
|
|
# dx[b,i] = sum_j g[b,j]*w[j,i]
|
|
# 注意 dw/db 按"累加"语义(调用前清零 Grad)
|
|
# ------------------------------------------------------------
|
|
def linear_bwd(w: Tensor | t.CPtr, x: Tensor | t.CPtr,
|
|
g: Tensor | t.CPtr, dw: Tensor | t.CPtr,
|
|
db: Tensor | t.CPtr, dx: Tensor | t.CPtr,
|
|
accum: int = 0):
|
|
if w is None or x is None or g is None:
|
|
return
|
|
if dw is None or db is None or dx is None:
|
|
return
|
|
m: sizet = x.Rows
|
|
k: sizet = x.Cols
|
|
n: sizet = w.Rows
|
|
wd: double | t.CPtr = w.Data
|
|
xd: double | t.CPtr = x.Data
|
|
gd: double | t.CPtr = g.Data
|
|
dwd: double | t.CPtr = dw.Data
|
|
dbd: double | t.CPtr = db.Data
|
|
dxd: double | t.CPtr = dx.Data
|
|
|
|
# dw/db/dx 清零(dx 恒清零: 它是本样本的输入梯度, 不跨样本累积)
|
|
# accum=1 时保留 dw/db 现值 —— 逐样本反向累积整批参数梯度的调用方用
|
|
if accum == 0:
|
|
string.memset(dw.Data, 0, n * k * sizet(8))
|
|
string.memset(db.Data, 0, n * sizet(8))
|
|
string.memset(dx.Data, 0, m * k * sizet(8))
|
|
|
|
b: sizet = 0
|
|
while b < m:
|
|
j: sizet = 0
|
|
while j < n:
|
|
gj: double = gd[b * n + j]
|
|
if gj != double(0.0):
|
|
dbd[j] = dbd[j] + gj
|
|
i: sizet = 0
|
|
while i < k:
|
|
dwd[j * k + i] = dwd[j * k + i] + gj * xd[b * k + i]
|
|
dxd[b * k + i] = dxd[b * k + i] + gj * wd[j * k + i]
|
|
i += 1
|
|
j += 1
|
|
b += 1
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# ReLU (simd_max4d 版本) 与反向
|
|
# ------------------------------------------------------------
|
|
def relu_fwd(x: Tensor | t.CPtr, out: Tensor | t.CPtr):
|
|
if x is None or out is None:
|
|
return
|
|
n: sizet = x.Rows * x.Cols
|
|
zero4: t.CPtr = stdlib.malloc(32)
|
|
if zero4 is None:
|
|
return
|
|
string.memset(zero4, 0, 32)
|
|
xd: double | t.CPtr = x.Data
|
|
od: double | t.CPtr = out.Data
|
|
i: sizet = 0
|
|
while i + 4 <= n:
|
|
vipersimd.simd_max4d(
|
|
t.CPtr(uint64(xd) + uint64(i) * 8),
|
|
zero4,
|
|
t.CPtr(uint64(od) + uint64(i) * 8))
|
|
i += 4
|
|
while i < n:
|
|
if xd[i] > double(0.0):
|
|
od[i] = xd[i]
|
|
else:
|
|
od[i] = double(0.0)
|
|
i += 1
|
|
stdlib.free(zero4)
|
|
|
|
|
|
def relu_backward(x: Tensor | t.CPtr, g: Tensor | t.CPtr,
|
|
gx: Tensor | t.CPtr):
|
|
"""gx = g * (x > 0)"""
|
|
if x is None or g is None or gx is None:
|
|
return
|
|
n: sizet = x.Rows * x.Cols
|
|
xd: double | t.CPtr = x.Data
|
|
gd: double | t.CPtr = g.Data
|
|
od: double | t.CPtr = gx.Data
|
|
i: sizet = 0
|
|
while i < n:
|
|
if xd[i] > double(0.0):
|
|
od[i] = gd[i]
|
|
else:
|
|
od[i] = double(0.0)
|
|
i += 1
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# softmax + CrossEntropy(融合, 数值稳定) —— MNIST 输出层
|
|
# logits: (m, C), labels: m 个 int (0..C-1)
|
|
# loss 写入 *loss, 梯度(已按 1/m 归一)写入 grad (m, C)
|
|
# ------------------------------------------------------------
|
|
def softmax_crossentropy(logits: Tensor | t.CPtr,
|
|
labels: int | t.CPtr,
|
|
n_labels: sizet,
|
|
loss: double | t.CPtr,
|
|
grad: Tensor | t.CPtr):
|
|
if logits is None or labels is None or grad is None:
|
|
return
|
|
if loss is None:
|
|
return
|
|
m: sizet = logits.Rows
|
|
c: sizet = logits.Cols
|
|
ld: double | t.CPtr = logits.Data
|
|
gd: double | t.CPtr = grad.Data
|
|
lb: int | t.CPtr = labels
|
|
total: double = double(0.0)
|
|
inv_m: double = double(1.0) / double(m)
|
|
b: sizet = 0
|
|
while b < m:
|
|
# 数值稳定 softmax
|
|
mx: double = ld[b * c]
|
|
j: sizet = 1
|
|
while j < c:
|
|
if ld[b * c + j] > mx:
|
|
mx = ld[b * c + j]
|
|
j += 1
|
|
s: double = double(0.0)
|
|
j = 0
|
|
while j < c:
|
|
s = s + math.exp(ld[b * c + j] - mx)
|
|
j += 1
|
|
y: int = lb[b]
|
|
# loss += -(logits[lab] - max - log(s))
|
|
total = total - (ld[b * c + sizet(y)] - mx - math.log(s))
|
|
# 梯度: (softmax - onehot) / m
|
|
j = 0
|
|
while j < c:
|
|
p: double = math.exp(ld[b * c + j] - mx) / s
|
|
o: double = double(0.0)
|
|
if sizet(y) == j:
|
|
o = double(1.0)
|
|
gd[b * c + j] = (p - o) * inv_m
|
|
j += 1
|
|
b += 1
|
|
loss[0] = total
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# 每行 argmax (分类预测)
|
|
# ------------------------------------------------------------
|
|
def row_argmax(x: Tensor | t.CPtr, pred: int | t.CPtr):
|
|
if x is None or pred is None:
|
|
return
|
|
m: sizet = x.Rows
|
|
c: sizet = x.Cols
|
|
d: double | t.CPtr = x.Data
|
|
i: sizet = 0
|
|
while i < m:
|
|
best: sizet = 0
|
|
bv: double = d[i * c]
|
|
j: sizet = 1
|
|
while j < c:
|
|
if d[i * c + j] > bv:
|
|
bv = d[i * c + j]
|
|
best = j
|
|
j += 1
|
|
pred[i] = int(best)
|
|
i += 1
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# SGD: data -= lr * grad(全量梯度) / 重球法: v = m*v+g, data -= lr*v
|
|
# ------------------------------------------------------------
|
|
def sgd_step(param: Tensor | t.CPtr, lr: double):
|
|
if param is None or param.Data is None or param.Grad is None:
|
|
return
|
|
d: double | t.CPtr = param.Data
|
|
g: double | t.CPtr = param.Grad
|
|
n: sizet = param.Rows * param.Cols
|
|
i: sizet = 0
|
|
while i < n:
|
|
d[i] = d[i] - lr * g[i]
|
|
i += 1
|
|
|
|
|
|
def sgd_step_momentum(param: Tensor | t.CPtr, vel: t.CPtr,
|
|
lr: double, momentum: double):
|
|
"""带动量(重球法)的 SGD: v = m*v + g ; w -= lr*v
|
|
|
|
vel 是与 param 同元素个数的 double 缓冲, 由调用方分配并清零(跨 step 持有);
|
|
vel 传 None 时退化为普通 sgd_step。动量是 CNN/MLP 训练收敛速度的关键 ——
|
|
纯 SGD 在固定 lr 下前几个 epoch 很慢且容易卡在低准确率。"""
|
|
if param is None or param.Data is None or param.Grad is None:
|
|
return
|
|
if vel is None:
|
|
sgd_step(param, lr)
|
|
return
|
|
d: double | t.CPtr = param.Data
|
|
g: double | t.CPtr = param.Grad
|
|
v: double | t.CPtr = vel
|
|
n: sizet = param.Rows * param.Cols
|
|
i: sizet = 0
|
|
while i < n:
|
|
vv: double = momentum * v[i] + g[i]
|
|
v[i] = vv
|
|
d[i] = d[i] - lr * vv
|
|
i += 1
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# 逐元素乘加 (调试/扩展用, vipersimd fma_array)
|
|
# a = a*b + c
|
|
# ------------------------------------------------------------
|
|
def fma_array_inplace(a: Tensor | t.CPtr, b: t.CPtr, c: t.CPtr):
|
|
if a is None or a.Data is None:
|
|
return
|
|
vipersimd.simd_fma_array(a.Data, b, c, a.Data, a.Rows * a.Cols)
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# Conv2d / MaxPool2d —— 轻量 CNN 用的原始缓冲算子
|
|
#
|
|
# 这些算子直接吃裸缓冲(不用 Tensor: torch.Tensor 只有 2D 形状),
|
|
# 布局一律 NCHW 连续 double; 卷积权重 [cout][cin][k][k]。
|
|
# 约定:
|
|
# *_fwd 覆盖写输出
|
|
# conv2d_bwd_dx 先清零再累加(逐样本输入梯度, 不跨样本累积)
|
|
# maxpool2d_bwd 同上
|
|
# conv2d_bwd_dw += 累加参数梯度(跨样本累积, 批次开始由调用方清零)
|
|
#
|
|
# SIMD: 三个卷积算子的内层都按"输出列"做 4 路 AVX2 double。
|
|
# 依据是 stride==1 时相邻输出列恰好对应相邻输入列, 故一列 4 个 double
|
|
# 在内存里连续, 可直接 vmovupd + vfmadd213pd, 权重标量 vbroadcastsd。
|
|
# 代价是需要一个"安全列区间": 4 条 lane 的输出列都在 [0,wout) 内, 且它们
|
|
# 触及的输入列 [ox-pad, ox+3+k-1-pad] 全在 [0,win) 内 —— 即
|
|
# ox ∈ [pad, win+pad-k-3]。区间外的列与 stride!=1 一律走原标量分支。
|
|
# (maxpool2d 不做 SIMD: k==stride==2 时 4 个窗口起点间隔 2 个 double,
|
|
# 不是连续地址, 无 gather 原语; 且它占总计算量不到 1%。)
|
|
# ------------------------------------------------------------
|
|
def _dptr(base: t.CPtr, elem: sizet) -> t.CPtr:
|
|
"""裸 double 缓冲的元素下标 → 字节地址
|
|
|
|
裸 t.CPtr 是 i8*, `p + n` 走的是字节偏移; 而 vipersimd 的 4 路原语按
|
|
"4 个 double = 32 字节"读写, 所以下标必须显式 ×8。
|
|
(与 linear_fwd 里的 t.CPtr(uint64(p) + i * 8) 同一写法。)"""
|
|
return t.CPtr(uint64(base) + uint64(elem) * 8)
|
|
|
|
|
|
def conv2d_fwd(x: t.CPtr, w: t.CPtr, b: t.CPtr, out: t.CPtr,
|
|
n: sizet, cin: sizet, hin: sizet, win: sizet,
|
|
cout: sizet, k: sizet, stride: sizet,
|
|
pad: sizet):
|
|
"""out[n][cout][hout][wout] = conv2d(x, w) + b
|
|
hout = (hin + 2*pad - k)/stride + 1(wout 同理, 整除向下取整)"""
|
|
if x is None or w is None or out is None:
|
|
return
|
|
hout: sizet = (hin + pad * 2 - k) / stride + 1
|
|
wout: sizet = (win + pad * 2 - k) / stride + 1
|
|
xd: double | t.CPtr = x
|
|
wd: double | t.CPtr = w
|
|
od: double | t.CPtr = out
|
|
kk: sizet = cin * k * k
|
|
# SIMD 安全区间(见文件头说明)
|
|
simd_on: int = 0
|
|
smax: sizet = 0
|
|
if stride == 1 and win + pad >= k + 3:
|
|
simd_on = 1
|
|
smax = win + pad - (k + 3)
|
|
wb: t.CPtr = None # 权重标量 × 4 路广播副本(kk 个标量)
|
|
t4: t.CPtr = None # 4 路累加器
|
|
if simd_on != 0:
|
|
wb = stdlib.malloc(kk * 32)
|
|
t4 = stdlib.malloc(32)
|
|
if wb is None or t4 is None:
|
|
if wb is not None:
|
|
stdlib.free(wb)
|
|
if t4 is not None:
|
|
stdlib.free(t4)
|
|
wb = None
|
|
t4 = None
|
|
simd_on = 0
|
|
wbd: double | t.CPtr = wb
|
|
t4d: double | t.CPtr = t4
|
|
|
|
bi: sizet = 0
|
|
while bi < n:
|
|
co: sizet = 0
|
|
while co < cout:
|
|
bias: double = double(0.0)
|
|
if b is not None:
|
|
bd: double | t.CPtr = b
|
|
bias = bd[co]
|
|
# 权重标量预广播: 每个 (ci,ky,kx) 一个 4 路副本, 本通道所有输出
|
|
# 行/列共用。否则 vbroadcastsd(4 次 store) 会跟着 4 路循环重复,
|
|
# 把 SIMD 的收益吃掉一半。
|
|
if simd_on != 0:
|
|
wb4: sizet = 0
|
|
while wb4 < kk:
|
|
wq: double = wd[co * kk + wb4]
|
|
wbd[wb4 * 4] = wq
|
|
wbd[wb4 * 4 + 1] = wq
|
|
wbd[wb4 * 4 + 2] = wq
|
|
wbd[wb4 * 4 + 3] = wq
|
|
wb4 += 1
|
|
oy: sizet = 0
|
|
while oy < hout:
|
|
oy0: sizet = oy * stride
|
|
ox: sizet = 0
|
|
while ox < wout:
|
|
if simd_on != 0 and ox >= pad and ox <= smax and ox + 4 <= wout:
|
|
# ---- 4 路: 一次算 4 个相邻输出列 ----
|
|
t4d[0] = bias
|
|
t4d[1] = bias
|
|
t4d[2] = bias
|
|
t4d[3] = bias
|
|
ci: sizet = 0
|
|
while ci < cin:
|
|
ky: sizet = 0
|
|
while ky < k:
|
|
iy: sizet = oy0 + ky
|
|
if iy >= pad:
|
|
yy: sizet = iy - pad
|
|
if yy < hin:
|
|
xrow: sizet = ((bi * cin + ci) * hin + yy) * win
|
|
kx: sizet = 0
|
|
while kx < k:
|
|
vipersimd.simd_fma4d(
|
|
_dptr(x, xrow + ox + kx - pad),
|
|
_dptr(wb, ((ci * k + ky) * k + kx) * 4),
|
|
t4, t4)
|
|
kx += 1
|
|
ky += 1
|
|
ci += 1
|
|
ob: sizet = ((bi * cout + co) * hout + oy) * wout + ox
|
|
od[ob] = t4d[0]
|
|
od[ob + 1] = t4d[1]
|
|
od[ob + 2] = t4d[2]
|
|
od[ob + 3] = t4d[3]
|
|
ox += 4
|
|
else:
|
|
# ---- 标量: 边界列 / stride != 1 ----
|
|
ox0: sizet = ox * stride
|
|
acc: double = bias
|
|
ci2: sizet = 0
|
|
while ci2 < cin:
|
|
ky2: sizet = 0
|
|
while ky2 < k:
|
|
iy2: sizet = oy0 + ky2
|
|
if iy2 >= pad:
|
|
yy2: sizet = iy2 - pad
|
|
if yy2 < hin:
|
|
kx2: sizet = 0
|
|
while kx2 < k:
|
|
ix2: sizet = ox0 + kx2
|
|
if ix2 >= pad:
|
|
xx2: sizet = ix2 - pad
|
|
if xx2 < win:
|
|
acc = acc + \
|
|
wd[co * kk + (ci2 * k + ky2) * k + kx2] * \
|
|
xd[((bi * cin + ci2) * hin + yy2) * win + xx2]
|
|
kx2 += 1
|
|
ky2 += 1
|
|
ci2 += 1
|
|
od[((bi * cout + co) * hout + oy) * wout + ox] = acc
|
|
ox += 1
|
|
oy += 1
|
|
co += 1
|
|
bi += 1
|
|
if wb is not None:
|
|
stdlib.free(wb)
|
|
if t4 is not None:
|
|
stdlib.free(t4)
|
|
|
|
|
|
def conv2d_bwd_dw(x: t.CPtr, dout: t.CPtr, dw: t.CPtr, db: t.CPtr,
|
|
n: sizet, cin: sizet, hin: sizet, win: sizet,
|
|
cout: sizet, k: sizet, stride: sizet,
|
|
pad: sizet):
|
|
"""累加: dw[co][ci][ky][kx] += sum dout*x ; db[co] += sum dout
|
|
(+= 语义, 调用方在批次开始前清零)
|
|
|
|
循环顺序为 (ci,ky,kx) 在外、(oy,ox) 在内: 每个权重元素只做一次
|
|
读-改-写; 内层按输出列 4 路 SIMD 累加, 最后 hsum 成标量。"""
|
|
if x is None or dout is None or dw is None:
|
|
return
|
|
hout: sizet = (hin + pad * 2 - k) / stride + 1
|
|
wout: sizet = (win + pad * 2 - k) / stride + 1
|
|
xd: double | t.CPtr = x
|
|
gd: double | t.CPtr = dout
|
|
dwd: double | t.CPtr = dw
|
|
dbd: double | t.CPtr = None
|
|
if db is not None:
|
|
dbd = db
|
|
kk: sizet = cin * k * k
|
|
simd_on: int = 0
|
|
smax: sizet = 0
|
|
if stride == 1 and win + pad >= k + 3:
|
|
simd_on = 1
|
|
smax = win + pad - (k + 3)
|
|
acc: t.CPtr = None
|
|
if simd_on != 0:
|
|
acc = stdlib.malloc(32)
|
|
if acc is None:
|
|
simd_on = 0
|
|
accd: double | t.CPtr = acc
|
|
|
|
bi: sizet = 0
|
|
while bi < n:
|
|
co: sizet = 0
|
|
while co < cout:
|
|
gp: sizet = (bi * cout + co) * hout * wout
|
|
# ---- db[co] += sum(g): g 的整个通道平面是连续的 ----
|
|
if dbd is not None:
|
|
tot0: double = double(0.0)
|
|
if simd_on != 0:
|
|
accd[0] = double(0.0)
|
|
accd[1] = double(0.0)
|
|
accd[2] = double(0.0)
|
|
accd[3] = double(0.0)
|
|
gcnt: sizet = hout * wout
|
|
gof: sizet = 0
|
|
while gof + 4 <= gcnt:
|
|
if simd_on != 0:
|
|
vipersimd.simd_add4d(_dptr(dout, gp + gof), acc, acc)
|
|
else:
|
|
tot0 = tot0 + gd[gp + gof]
|
|
tot0 = tot0 + gd[gp + gof + 1]
|
|
tot0 = tot0 + gd[gp + gof + 2]
|
|
tot0 = tot0 + gd[gp + gof + 3]
|
|
gof += 4
|
|
while gof < gcnt:
|
|
tot0 = tot0 + gd[gp + gof]
|
|
gof += 1
|
|
if simd_on != 0:
|
|
vipersimd.simd_hsum4d(acc, acc)
|
|
tot0 = tot0 + accd[0]
|
|
dbd[co] = dbd[co] + tot0
|
|
# ---- dw[co][ci][ky][kx] ----
|
|
ci: sizet = 0
|
|
while ci < cin:
|
|
ky: sizet = 0
|
|
while ky < k:
|
|
kx: sizet = 0
|
|
while kx < k:
|
|
tot: double = double(0.0)
|
|
if simd_on != 0:
|
|
accd[0] = double(0.0)
|
|
accd[1] = double(0.0)
|
|
accd[2] = double(0.0)
|
|
accd[3] = double(0.0)
|
|
oy: sizet = 0
|
|
while oy < hout:
|
|
iy: sizet = oy * stride + ky
|
|
if iy >= pad:
|
|
yy: sizet = iy - pad
|
|
if yy < hin:
|
|
xrow: sizet = ((bi * cin + ci) * hin + yy) * win
|
|
goff: sizet = gp + oy * wout
|
|
ox: sizet = 0
|
|
while ox < wout:
|
|
if simd_on != 0 and ox >= pad and ox <= smax and ox + 4 <= wout:
|
|
vipersimd.simd_fma4d(
|
|
_dptr(dout, goff + ox),
|
|
_dptr(x, xrow + ox + kx - pad),
|
|
acc, acc)
|
|
ox += 4
|
|
else:
|
|
ix: sizet = ox * stride + kx
|
|
if ix >= pad:
|
|
xx: sizet = ix - pad
|
|
if xx < win:
|
|
tot = tot + gd[goff + ox] * \
|
|
xd[xrow + xx]
|
|
ox += 1
|
|
oy += 1
|
|
if simd_on != 0:
|
|
vipersimd.simd_hsum4d(acc, acc)
|
|
tot = tot + accd[0]
|
|
wj: sizet = co * kk + (ci * k + ky) * k + kx
|
|
dwd[wj] = dwd[wj] + tot
|
|
kx += 1
|
|
ky += 1
|
|
ci += 1
|
|
co += 1
|
|
bi += 1
|
|
if acc is not None:
|
|
stdlib.free(acc)
|
|
|
|
|
|
def conv2d_bwd_dx(w: t.CPtr, dout: t.CPtr, dx: t.CPtr,
|
|
n: sizet, cin: sizet, hin: sizet, win: sizet,
|
|
cout: sizet, k: sizet, stride: sizet,
|
|
pad: sizet):
|
|
"""dx 先清零再累加(逐样本输入梯度): dx[ci][iy][ix] += sum dout*w
|
|
|
|
stride==1 时内层按输出列 4 路 SIMD: 4 条 lane 对应的 dx 单元也是连续的,
|
|
故可直接对 dx 做 in-place fma(g4 * w_broadcast + dx4)。"""
|
|
if w is None or dout is None or dx is None:
|
|
return
|
|
hout: sizet = (hin + pad * 2 - k) / stride + 1
|
|
wout: sizet = (win + pad * 2 - k) / stride + 1
|
|
wd: double | t.CPtr = w
|
|
gd: double | t.CPtr = dout
|
|
dxd: double | t.CPtr = dx
|
|
string.memset(dx, 0, n * cin * hin * win * sizet(8))
|
|
kk: sizet = cin * k * k
|
|
simd_on: int = 0
|
|
smax: sizet = 0
|
|
if stride == 1 and win + pad >= k + 3:
|
|
simd_on = 1
|
|
smax = win + pad - (k + 3)
|
|
wb: t.CPtr = None # 整个权重张量的 4 路广播副本(cout*kk 个标量)
|
|
if simd_on != 0:
|
|
wb = stdlib.malloc(cout * kk * 32)
|
|
if wb is None:
|
|
simd_on = 0
|
|
wbd: double | t.CPtr = wb
|
|
if simd_on != 0:
|
|
# 权重在一次调用内不变 → 整体预广播一次(注意别在 4 路循环里重复
|
|
# vbroadcastsd: 那是 4 次 store, 会把 SIMD 的收益吃掉一半)
|
|
wjn: sizet = cout * kk
|
|
wj4: sizet = 0
|
|
while wj4 < wjn:
|
|
wq2: double = wd[wj4]
|
|
wbd[wj4 * 4] = wq2
|
|
wbd[wj4 * 4 + 1] = wq2
|
|
wbd[wj4 * 4 + 2] = wq2
|
|
wbd[wj4 * 4 + 3] = wq2
|
|
wj4 += 1
|
|
|
|
bi: sizet = 0
|
|
while bi < n:
|
|
ci: sizet = 0
|
|
while ci < cin:
|
|
ky: sizet = 0
|
|
while ky < k:
|
|
oy: sizet = 0
|
|
while oy < hout:
|
|
iy: sizet = oy * stride + ky
|
|
if iy >= pad:
|
|
yy: sizet = iy - pad
|
|
if yy < hin:
|
|
dxrow: sizet = ((bi * cin + ci) * hin + yy) * win
|
|
ox: sizet = 0
|
|
while ox < wout:
|
|
if simd_on != 0 and ox >= pad and ox <= smax and ox + 4 <= wout:
|
|
co: sizet = 0
|
|
while co < cout:
|
|
kx: sizet = 0
|
|
while kx < k:
|
|
goff: sizet = ((bi * cout + co) * hout + oy) * wout + ox
|
|
doff: sizet = dxrow + ox + kx - pad
|
|
vipersimd.simd_fma4d(
|
|
_dptr(dout, goff),
|
|
_dptr(wb, (co * kk + (ci * k + ky) * k + kx) * 4),
|
|
_dptr(dx, doff), _dptr(dx, doff))
|
|
kx += 1
|
|
co += 1
|
|
ox += 4
|
|
else:
|
|
co2: sizet = 0
|
|
while co2 < cout:
|
|
kx2: sizet = 0
|
|
while kx2 < k:
|
|
ix: sizet = ox * stride + kx2
|
|
if ix >= pad:
|
|
xx: sizet = ix - pad
|
|
if xx < win:
|
|
xi: sizet = dxrow + xx
|
|
gi: sizet = ((bi * cout + co2) * hout + oy) * wout + ox
|
|
dxd[xi] = dxd[xi] + \
|
|
gd[gi] * wd[co2 * kk + (ci * k + ky) * k + kx2]
|
|
kx2 += 1
|
|
co2 += 1
|
|
ox += 1
|
|
oy += 1
|
|
ky += 1
|
|
ci += 1
|
|
bi += 1
|
|
if wb is not None:
|
|
stdlib.free(wb)
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# MaxPool2d (k x k, stride 通常 = k)
|
|
# ------------------------------------------------------------
|
|
def maxpool2d_fwd(x: t.CPtr, out: t.CPtr, argmax: t.CPtr,
|
|
n: sizet, c: sizet, hin: sizet, win: sizet,
|
|
k: sizet, stride: sizet):
|
|
"""argmax: int 数组 [n*c*hout*wout], 记窗口内最大值在本通道内的下标(yy*win+xx)"""
|
|
if x is None or out is None or argmax is None:
|
|
return
|
|
hout: sizet = (hin - k) / stride + 1
|
|
wout: sizet = (win - k) / stride + 1
|
|
xd: double | t.CPtr = x
|
|
od: double | t.CPtr = out
|
|
am: int | t.CPtr = argmax
|
|
bi: sizet = 0
|
|
while bi < n:
|
|
ch: sizet = 0
|
|
while ch < c:
|
|
plano: sizet = (bi * c + ch) * hin * win
|
|
oy: sizet = 0
|
|
while oy < hout:
|
|
oy0: sizet = oy * stride
|
|
ox: sizet = 0
|
|
while ox < wout:
|
|
ox0: sizet = ox * stride
|
|
best: double = double(0.0)
|
|
bidx: int = 0
|
|
first: int = 1
|
|
ky: sizet = 0
|
|
while ky < k:
|
|
yy: sizet = oy0 + ky
|
|
kx: sizet = 0
|
|
while kx < k:
|
|
xx: sizet = ox0 + kx
|
|
v: double = xd[plano + yy * win + xx]
|
|
if first != 0 or v > best:
|
|
best = v
|
|
bidx = int(yy * win + xx)
|
|
first = 0
|
|
kx += 1
|
|
ky += 1
|
|
oi: sizet = (bi * c + ch) * hout * wout + oy * wout + ox
|
|
od[oi] = best
|
|
am[oi] = bidx
|
|
ox += 1
|
|
oy += 1
|
|
ch += 1
|
|
bi += 1
|
|
|
|
|
|
def maxpool2d_bwd(dout: t.CPtr, argmax: t.CPtr, dx: t.CPtr,
|
|
n: sizet, c: sizet, hin: sizet, win: sizet,
|
|
k: sizet, stride: sizet):
|
|
"""dx 先清零再散射累加: dx[argmax] += dout"""
|
|
if dout is None or argmax is None or dx is None:
|
|
return
|
|
hout: sizet = (hin - k) / stride + 1
|
|
wout: sizet = (win - k) / stride + 1
|
|
gd: double | t.CPtr = dout
|
|
am: int | t.CPtr = argmax
|
|
dxd: double | t.CPtr = dx
|
|
string.memset(dx, 0, n * c * hin * win * sizet(8))
|
|
bi: sizet = 0
|
|
while bi < n:
|
|
ch: sizet = 0
|
|
while ch < c:
|
|
plano: sizet = (bi * c + ch) * hin * win
|
|
oy: sizet = 0
|
|
while oy < hout:
|
|
ox: sizet = 0
|
|
while ox < wout:
|
|
oi: sizet = (bi * c + ch) * hout * wout + oy * wout + ox
|
|
tgt: sizet = plano + sizet(am[oi])
|
|
dxd[tgt] = dxd[tgt] + gd[oi]
|
|
ox += 1
|
|
oy += 1
|
|
ch += 1
|
|
bi += 1
|