GVSDS/Viper_VRTE
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(torch): vector_add_inplace 正确调用 simd_add_array

Browse Source 
simd_add_array(a, b, out, n) 是整段批量函数，原先只传了 3 个参数（缺 count），且在外层每 4 个元素调用一次、每次从下标 0 重新开始。改为单次调用 simd_add_array(a, b, a, n)，尾部用 n-(n&3) 起做标量补齐；顺带删掉无用的 malloc(8) 临时块。注：本次提交同时带上了该文件既有的 t.CPtr → rawptr 迁移改动。
This commit is contained in:
TermiNexus committed 2026-10-17 18:47:57 +08:00
1 parent 1b21d58be3
commit 7556264d22
1 file changed
+36 -48
Download Patch File Download Diff File Expand all files Collapse all files
torch/__init__.vp
+36 -48
View File
@@ -16,11 +16,8 @@
# * Memory: all Tensors and their data/grad are allocated from a caller-provided memhub.MemBuddy
# (consistent with the existing includes convention, pool.alloc avoids manual free).
# ============================================================
from stdint import *
import t, c
import stdlib
import string
import stdio
import math
import vipersimd
import memhub
@@ -30,8 +27,8 @@ 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)
Data: rawptr # rows*cols*8 contiguous bytes of double
Grad: rawptr # same-shape gradient (may be None)
Rows: sizet
Cols: sizet
@@ -39,9 +36,9 @@ novirtual class Tensor:
define TENSOR_SIZE = 32
def _zeros_block(pool: memhub.MemBuddy*, bytes_n: sizet) -> t.CPtr:
def _zeros_block(pool: memhub.MemBuddy*, bytes_n: sizet) -> rawptr:
"""Allocate and zero a block of memory"""
p: t.CPtr = pool.alloc(bytes_n)
p: rawptr = pool.alloc(bytes_n)
if p is not None:
string.memset(p, 0, bytes_n)
return p
@@ -67,7 +64,7 @@ def zeros(pool: memhub.MemBuddy*, rows: sizet,
return t0
def from_raw(pool: memhub.MemBuddy*, data: t.CPtr,
def from_raw(pool: memhub.MemBuddy*, data: rawptr,
rows: sizet, cols: sizet,
with_grad: int = 0) -> Tensor*:
"""torch.from_raw: wrap existing memory as a Tensor (no copy, not owned)"""
@@ -89,13 +86,13 @@ def tensor_free(t0: Tensor*):
pass
def tensor_data(t0: Tensor*) -> t.CPtr:
def tensor_data(t0: Tensor*) -> rawptr:
if t0 is None:
return None
return t0.Data
def tensor_grad(t0: Tensor*) -> t.CPtr:
def tensor_grad(t0: Tensor*) -> rawptr:
if t0 is None:
return None
return t0.Grad
@@ -174,17 +171,14 @@ def kaiming_uniform_fill(w: Tensor*, fan_in: sizet,
# ------------------------------------------------------------
# Basic element-wise ops (using vipersimd 4-wide instructions)
# ------------------------------------------------------------
def vector_mul_scalar4(src: t.CPtr, scalar_ptr: t.CPtr, dst: t.CPtr,
def vector_mul_scalar4(src: rawptr, scalar_ptr: rawptr, dst: rawptr,
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))
vipersimd.simd_mul_scalar4d(rawptr(uint64(src) + uint64(i) * 8), scalar_ptr, rawptr(uint64(dst) + uint64(i) * 8))
i += 4
sd: double* = src
dd: double* = dst
@@ -196,26 +190,21 @@ def vector_mul_scalar4(src: t.CPtr, scalar_ptr: t.CPtr, dst: t.CPtr,
i += 1
def vector_add_inplace(a: t.CPtr, b: t.CPtr, n: sizet):
def vector_add_inplace(a: rawptr, b: rawptr, 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
# simd_add_array is a *batch* over the whole range: it walks every full 4-double vector
# (floor(n/4)*4 elements) from the given pointers, so it takes (a, b, out, n) and must be
# called once. The previous per-4 loop passed only 3 arguments (missing the count) and
# restarted at index 0 on every iteration.
vipersimd.simd_add_array(da, db, da, n)
i: sizet = n - (n & 3)
while i < n:
da[i] = da[i] + db[i]
i += 1
stdlib.free(tmp)
# ------------------------------------------------------------
@@ -233,7 +222,7 @@ def linear_fwd(w: Tensor*, b: Tensor*,
bd: double* = b.Data
xd: double* = x.Data
od: double* = out.Data
acc_buf: t.CPtr = stdlib.malloc(8)
acc_buf: rawptr = stdlib.malloc(8)
if acc_buf is None:
return
accp: double* = acc_buf
@@ -243,8 +232,8 @@ def linear_fwd(w: Tensor*, b: Tensor*,
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),
rawptr(uint64(xd) + uint64(i * k + kk) * 8),
rawptr(uint64(wd) + uint64(j * k + kk) * 8),
acc_buf)
acc = acc + accp[0]
kk += 4
@@ -303,7 +292,7 @@ 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)
zero4: rawptr = stdlib.malloc(32)
if zero4 is None:
return
string.memset(zero4, 0, 32)
@@ -312,9 +301,9 @@ def relu_fwd(x: Tensor*, out: Tensor*):
i: sizet = 0
while i + 4 <= n:
vipersimd.simd_max4d(
t.CPtr(uint64(xd) + uint64(i) * 8),
rawptr(uint64(xd) + uint64(i) * 8),
zero4,
t.CPtr(uint64(od) + uint64(i) * 8))
rawptr(uint64(od) + uint64(i) * 8))
i += 4
while i < n:
if xd[i] > double(0.0):
@@ -424,8 +413,7 @@ def sgd_step(param: Tensor*, lr: double):
d[i] = d[i] - lr * g[i]
def sgd_step_momentum(param: Tensor*, vel: t.CPtr,
lr: double, momentum: double):
def sgd_step_momentum(param: Tensor*, vel: rawptr, 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:
@@ -447,7 +435,7 @@ def sgd_step_momentum(param: Tensor*, vel: t.CPtr,
# 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):
def fma_array_inplace(a: Tensor*, b: rawptr, c: rawptr):
if a is None or a.Data is None:
return
vipersimd.simd_fma_array(a.Data, b, c, a.Data, a.Rows * a.Cols)
@@ -458,13 +446,13 @@ def fma_array_inplace(a: Tensor*, b: t.CPtr, c: t.CPtr):
# 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:
def _dptr(base: rawptr, elem: sizet) -> rawptr:
"""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)
Raw rawptr 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 rawptr(uint64(base) + uint64(elem) * 8)
def conv2d_fwd(x: t.CPtr, w: t.CPtr, b: t.CPtr, out: t.CPtr,
def conv2d_fwd(x: rawptr, w: rawptr, b: rawptr, out: rawptr,
n: sizet, cin: sizet, hin: sizet, win: sizet,
cout: sizet, k: sizet, stride: sizet,
pad: sizet):
@@ -484,8 +472,8 @@ def conv2d_fwd(x: t.CPtr, w: t.CPtr, b: t.CPtr, out: t.CPtr,
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
wb: rawptr = None # 4-lane broadcast copy of weight scalars (kk scalars)
t4: rawptr = None # 4-lane accumulator
if simd_on != 0:
wb = stdlib.malloc(kk * 32)
t4 = stdlib.malloc(32)
@@ -569,7 +557,7 @@ def conv2d_fwd(x: t.CPtr, w: t.CPtr, b: t.CPtr, out: t.CPtr,
stdlib.free(t4)
def conv2d_bwd_dw(x: t.CPtr, dout: t.CPtr, dw: t.CPtr, db: t.CPtr,
def conv2d_bwd_dw(x: rawptr, dout: rawptr, dw: rawptr, db: rawptr,
n: sizet, cin: sizet, hin: sizet, win: sizet,
cout: sizet, k: sizet, stride: sizet,
pad: sizet):
@@ -591,7 +579,7 @@ def conv2d_bwd_dw(x: t.CPtr, dout: t.CPtr, dw: t.CPtr, db: t.CPtr,
if stride == 1 and win + pad >= k + 3:
simd_on = 1
smax = win + pad - (k + 3)
acc: t.CPtr = None
acc: rawptr = None
if simd_on != 0:
acc = stdlib.malloc(32)
if acc is None:
@@ -669,7 +657,7 @@ def conv2d_bwd_dw(x: t.CPtr, dout: t.CPtr, dw: t.CPtr, db: t.CPtr,
stdlib.free(acc)
def conv2d_bwd_dx(w: t.CPtr, dout: t.CPtr, dx: t.CPtr,
def conv2d_bwd_dx(w: rawptr, dout: rawptr, dx: rawptr,
n: sizet, cin: sizet, hin: sizet, win: sizet,
cout: sizet, k: sizet, stride: sizet,
pad: sizet):
@@ -689,7 +677,7 @@ def conv2d_bwd_dx(w: t.CPtr, dout: t.CPtr, dx: t.CPtr,
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)
wb: rawptr = 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:
@@ -745,7 +733,7 @@ def conv2d_bwd_dx(w: t.CPtr, dout: t.CPtr, dx: t.CPtr,
# ------------------------------------------------------------
# MaxPool2d (k x k, stride usually = k)
# ------------------------------------------------------------
def maxpool2d_fwd(x: t.CPtr, out: t.CPtr, argmax: t.CPtr,
def maxpool2d_fwd(x: rawptr, out: rawptr, argmax: rawptr,
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)"""
@@ -780,7 +768,7 @@ def maxpool2d_fwd(x: t.CPtr, out: t.CPtr, argmax: t.CPtr,
am[oi] = bidx
def maxpool2d_bwd(dout: t.CPtr, argmax: t.CPtr, dx: t.CPtr,
def maxpool2d_bwd(dout: rawptr, argmax: rawptr, dx: rawptr,
n: sizet, c: sizet, hin: sizet, win: sizet,
k: sizet, stride: sizet):
"""dx is zeroed then scatter-accumulated: dx[argmax] += dout"""