diff --git a/torch/__init__.vp b/torch/__init__.vp index da93a86..70e1c41 100644 --- a/torch/__init__.vp +++ b/torch/__init__.vp @@ -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"""