Viper_VRTE/torch/__init__.vp

# -*- 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 contiguous bytes of double
    Grad: t.CPtr          # same-shape gradient (may be None)
    Rows: sizet
    Cols: sizet

# Fields Data/Grad/Rows/Cols = 32 bytes (no vtable)
define TENSOR_SIZE = 32


def _zeros_block(pool: memhub.MemBuddy*, bytes_n: sizet) -> t.CPtr:
    """Allocate and zero a block of memory"""
    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*:
    """torch.zeros: create a rows x cols all-zero Tensor, optionally allocating a Grad buffer"""
    if pool is None:
        return None
    t0: 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*:
    """torch.from_raw: wrap existing memory as a Tensor (no copy, not owned)"""
    if pool is None:
        return None
    t0: 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*):
    """Free the Tensor shell (data is reclaimed by the MemBuddy pool; no manual free needed)"""
    pass


def tensor_data(t0: Tensor*) -> t.CPtr:
    if t0 is None:
        return None
    return t0.Data


def tensor_grad(t0: Tensor*) -> t.CPtr:
    if t0 is None:
        return None
    return t0.Grad


def tensor_rows(t0: Tensor*) -> sizet:
    if t0 is None:
        return 0
    return t0.Rows


def tensor_cols(t0: Tensor*) -> sizet:
    if t0 is None:
        return 0
    return t0.Cols


def fill(t0: Tensor*, val: double):
    """torch.fill: set all elements to val"""
    if t0 is None or t0.Data is None:
        return
    d: double* = t0.Data
    n: sizet = t0.Rows * t0.Cols
    for i in range(n):
        d[i] = val


def zero_grad(t0: Tensor*):
    """torch.zero_grad: zero the Grad buffer"""
    if t0 is None or t0.Grad is None:
        return
    string.memset(t0.Grad, 0, t0.Rows * t0.Cols * sizet(8))


# ------------------------------------------------------------
# Deterministic pseudo-random (LCG, minstd) and Kaiming uniform init
# ------------------------------------------------------------
_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) uniform random"""
    global _rng_state
    _rng_state = math.fmod(_rng_state * LCG_A, LCG_M)
    return _rng_state / LCG_M


def uniform_fill(t0: Tensor*, lo: double, hi: double):
    """Fill with U(lo, hi)"""
    if t0 is None or t0.Data is None:
        return
    d: double* = t0.Data
    n: sizet = t0.Rows * t0.Cols
    span: double = hi - lo
    for i in range(n):
        d[i] = lo + _next_rand() * span


def kaiming_uniform_fill(w: Tensor*, fan_in: sizet,
                         fan_out: sizet):
    """PyTorch Linear default init: U(-sqrt(6/(fan_in+fan_out)), +...)"""
    bound: double = math.sqrt(
        double(6.0) / double(fan_in + fan_out))
    uniform_fill(w, -bound, bound)


# ------------------------------------------------------------
# Basic element-wise ops (using vipersimd 4-wide instructions)
# ------------------------------------------------------------
def vector_mul_scalar4(src: t.CPtr, scalar_ptr: t.CPtr, dst: t.CPtr,
                       n: sizet):
    """dst = src * scalar (4-wide simd_mul_scalar4d + scalar tail)"""
    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* = src
    dd: double* = dst
    # Scalar tail
    sp: double* = 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-wide simd_add_array + scalar tail)"""
    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* = a
    db: double* = b
    while i < n:
        da[i] = da[i] + db[i]
        i += 1
    stdlib.free(tmp)


# ------------------------------------------------------------
# Linear layer: forward = x @ W^T + b  (W: out x in, x: m x in, out: m x out)
# Inner product in k%4 chunks via vipersimd.simd_dot4d -- the MNIST compute hotspot
# ------------------------------------------------------------
def linear_fwd(w: Tensor*, b: Tensor*,
               x: Tensor*, out: Tensor*):
    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* = w.Data
    bd: double* = b.Data
    xd: double* = x.Data
    od: double* = out.Data
    acc_buf: t.CPtr = stdlib.malloc(8)
    if acc_buf is None:
        return
    accp: double* = acc_buf
    for i in range(m):
        for j in range(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
    stdlib.free(acc_buf)


# ------------------------------------------------------------
# Linear backward: inputs 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]
# Note: dw/db use accumulate semantics (zero Grad before the call)
# ------------------------------------------------------------
def linear_bwd(w: Tensor*, x: Tensor*,
               g: Tensor*, dw: Tensor*,
               db: Tensor*, dx: Tensor*,
               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* = w.Data
    xd: double* = x.Data
    gd: double* = g.Data
    dwd: double* = dw.Data
    dbd: double* = db.Data
    dxd: double* = dx.Data

    # Zero dw/db/dx (dx is always zeroed: it is this sample's input gradient, not accumulated across samples)
    # accum=1 keeps dw/db -- for callers accumulating whole-batch parameter gradients sample by sample
    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))

    for b in range(m):
        for j in range(n):
            gj: double = gd[b * n + j]
            if gj != double(0.0):
                dbd[j] = dbd[j] + gj
                for i in range(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]


# ------------------------------------------------------------
# ReLU (simd_max4d version) and backward
# ------------------------------------------------------------
def relu_fwd(x: Tensor*, out: Tensor*):
    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* = x.Data
    od: double* = 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*, g: Tensor*,
                  gx: Tensor*):
    """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* = x.Data
    gd: double* = g.Data
    od: double* = gx.Data
    for i in range(n):
        if xd[i] > double(0.0):
            od[i] = gd[i]
        else:
            od[i] = double(0.0)


# ------------------------------------------------------------
# softmax + CrossEntropy (fused, numerically stable) -- MNIST output layer
# logits: (m, C), labels: m ints (0..C-1)
# loss written to *loss, gradient (normalized by 1/m) written to grad (m, C)
# ------------------------------------------------------------
def softmax_crossentropy(logits: Tensor*,
                         labels: int*,
                         n_labels: sizet,
                         loss: double*,
                         grad: Tensor*):
    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* = logits.Data
    gd: double* = grad.Data
    lb: int* = labels
    total: double = double(0.0)
    inv_m: double = double(1.0) / double(m)
    for b in range(m):
        # Numerically stable 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))
        # Gradient: (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
    loss[0] = total


# ------------------------------------------------------------
# Per-row argmax (classification prediction)
# ------------------------------------------------------------
def row_argmax(x: Tensor*, pred: int*):
    if x is None or pred is None:
        return
    m: sizet = x.Rows
    c: sizet = x.Cols
    d: double* = x.Data
    for i in range(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)


# ------------------------------------------------------------
# SGD: data -= lr * grad (full gradient)  /  momentum: v = m*v+g, data -= lr*v
# ------------------------------------------------------------
def sgd_step(param: Tensor*, lr: double):
    if param is None or param.Data is None or param.Grad is None:
        return
    d: double* = param.Data
    g: double* = param.Grad
    n: sizet = param.Rows * param.Cols
    for i in range(n):
        d[i] = d[i] - lr * g[i]


def sgd_step_momentum(param: Tensor*, vel: t.CPtr,
                      lr: double, momentum: double):
    """SGD with momentum: v = m*v + g ; w -= lr*v
    vel is a caller-allocated, zeroed double buffer with the same element count as param (held across steps); pass None to fall back to plain sgd_step."""
    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* = param.Data
    g: double* = param.Grad
    v: double* = vel
    n: sizet = param.Rows * param.Cols
    for i in range(n):
        vv: double = momentum * v[i] + g[i]
        v[i] = vv
        d[i] = d[i] - lr * vv


# ------------------------------------------------------------
# Element-wise multiply-add (for debug/extension, vipersimd fma_array)
# a = a*b + c
# ------------------------------------------------------------
def fma_array_inplace(a: Tensor*, 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 raw-buffer kernels for a lightweight CNN (NCHW contiguous double, weights [cout][cin][k][k]).
# Conventions: *_fwd overwrites output; conv2d_bwd_dx/maxpool2d_bwd zero then accumulate (per-sample dx); conv2d_bwd_dw accumulates params (+=, caller zeroes per batch).
# SIMD inner loop over output columns, safe range ox in [pad, win+pad-k-3] when stride==1 (maxpool2d is scalar only).
# ------------------------------------------------------------
def _dptr(base: t.CPtr, elem: sizet) -> t.CPtr:
    """Convert an element index into a byte address for a raw double buffer.
    Raw t.CPtr is i8* and `p + n` is a byte offset, so the index must be multiplied by 8 for vipersimd 4-wide (32-byte) access."""
    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 (same for wout; integer division floors)."""
    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* = x
    wd: double* = w
    od: double* = out
    kk: sizet = cin * k * k
    # SIMD safe range (see file-header notes)
    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-lane broadcast copy of weight scalars (kk scalars)
    t4: t.CPtr = None            # 4-lane accumulator
    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* = wb
    t4d: double* = t4

    for bi in range(n):
        for co in range(cout):
            bias: double = double(0.0)
            if b is not None:
                bd: double* = b
                bias = bd[co]
            # Pre-broadcast weight scalars: one 4-lane copy per (ci,ky,kx), shared by all output rows/cols of this channel.
            # Otherwise vbroadcastsd (4 stores) repeats in the 4-lane loop and halves the SIMD gain.
            if simd_on != 0:
                for wb4 in range(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
            for oy in range(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-lane: compute 4 adjacent output columns at once ----
                        t4d[0] = bias
                        t4d[1] = bias
                        t4d[2] = bias
                        t4d[3] = bias
                        for ci in range(cin):
                            for ky in range(k):
                                iy: sizet = oy0 + ky
                                if iy >= pad:
                                    yy: sizet = iy - pad
                                    if yy < hin:
                                        xrow: sizet = ((bi * cin + ci) * hin + yy) * win
                                        for kx in range(k):
                                            vipersimd.simd_fma4d(
                                                _dptr(x, xrow + ox + kx - pad),
                                                _dptr(wb, ((ci * k + ky) * k + kx) * 4),
                                                t4, t4)
                        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:
                        # ---- scalar: boundary columns / stride != 1 ----
                        ox0: sizet = ox * stride
                        acc: double = bias
                        for ci2 in range(cin):
                            for ky2 in range(k):
                                iy2: sizet = oy0 + ky2
                                if iy2 >= pad:
                                    yy2: sizet = iy2 - pad
                                    if yy2 < hin:
                                        for kx2 in range(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]
                        od[((bi * cout + co) * hout + oy) * wout + ox] = acc
                        ox += 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):
    """Accumulate: dw[co][ci][ky][kx] += sum dout*x ; db[co] += sum dout (+= semantics, caller zeroes before the batch).
    Loop order (ci,ky,kx) outer, (oy,ox) inner: one read-modify-write per weight; inner loop 4-lane SIMD, then hsum to scalar."""
    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* = x
    gd: double* = dout
    dwd: double* = dw
    dbd: double* = 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* = acc

    for bi in range(n):
        for co in range(cout):
            gp: sizet = (bi * cout + co) * hout * wout
            # ---- db[co] += sum(g): the whole channel plane of g is contiguous ----
            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] ----
            for ci in range(cin):
                for ky in range(k):
                    for kx in range(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)
                        for oy in range(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
                        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
    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 is zeroed then accumulated (per-sample input gradient): dx[ci][iy][ix] += sum dout*w.
    With stride==1 the inner loop is 4-lane SIMD; the lanes map to contiguous dx cells, so dx can be updated in place with fma(g4 * w4 + 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* = w
    gd: double* = dout
    dxd: double* = 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-lane broadcast copy of the whole weight tensor (cout*kk scalars)
    if simd_on != 0:
        wb = stdlib.malloc(cout * kk * 32)
        if wb is None:
            simd_on = 0
    wbd: double* = wb
    if simd_on != 0:
        # Weights are constant within one call, so pre-broadcast once (do not repeat vbroadcastsd in the 4-lane loop: 4 stores halves the SIMD gain)
        wjn: sizet = cout * kk
        for wj4 in range(wjn):
            wq2: double = wd[wj4]
            wbd[wj4 * 4] = wq2
            wbd[wj4 * 4 + 1] = wq2
            wbd[wj4 * 4 + 2] = wq2
            wbd[wj4 * 4 + 3] = wq2

    for bi in range(n):
        for ci in range(cin):
            for ky in range(k):
                for oy in range(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:
                                    for co in range(cout):
                                        for kx in range(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))
                                    ox += 4
                                else:
                                    for co2 in range(cout):
                                        for kx2 in range(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]
                                    ox += 1
    if wb is not None:
        stdlib.free(wb)


# ------------------------------------------------------------
# MaxPool2d (k x k, stride usually = 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 array [n*c*hout*wout] recording the max index within the channel (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* = x
    od: double* = out
    am: int* = argmax
    for bi in range(n):
        for ch in range(c):
            plano: sizet = (bi * c + ch) * hin * win
            for oy in range(hout):
                oy0: sizet = oy * stride
                for ox in range(wout):
                    ox0: sizet = ox * stride
                    best: double = double(0.0)
                    bidx: int = 0
                    first: int = 1
                    for ky in range(k):
                        yy: sizet = oy0 + ky
                        for kx in range(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
                    oi: sizet = (bi * c + ch) * hout * wout + oy * wout + ox
                    od[oi] = best
                    am[oi] = bidx


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 is zeroed then scatter-accumulated: 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* = dout
    am: int* = argmax
    dxd: double* = dx
    string.memset(dx, 0, n * c * hin * win * sizet(8))
    for bi in range(n):
        for ch in range(c):
            plano: sizet = (bi * c + ch) * hin * win
            for oy in range(hout):
                for ox in range(wout):
                    oi: sizet = (bi * c + ch) * hout * wout + oy * wout + ox
                    tgt: sizet = plano + sizet(am[oi])
                    dxd[tgt] = dxd[tgt] + gd[oi]