Files
TransPyC/Test/SimdTest/App/main.py
2026-06-16 16:09:42 +08:00

518 lines
16 KiB
Python

from stdint import *
import vipersimd
import numpy
from stdio import printf
import stdlib
from stdlib import calloc
import t, c
import mpool
# 全局内存池
arena: t.CUInt8T | t.CPtr = None
pool: mpool.MPool | t.CPtr = None
# 外部函数: clock() 用于性能计时
def clock() -> t.CLong | t.CExtern: pass
CLOCK_PER_SEC: t.CDefine = 1000
test_passed: t.CInt = 0
test_failed: t.CInt = 0
def check(name: str, condition: t.CInt, detail: str):
global test_passed, test_failed
if condition:
test_passed += 1
printf(" [PASS] %s\n", name)
else:
test_failed += 1
printf(" [FAIL] %s -- %s\n", name, detail if detail else "")
def section(title: str):
printf("\n=== %s ===\n", title)
# ============================================================
# 1. AVX2 double4 正确性测试
# ============================================================
def test_avx2_add_sub_mul_div():
section("AVX2 add/sub/mul/div 4d")
a: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
a.data[0] = t.CDouble(1.0)
a.data[1] = t.CDouble(2.0)
a.data[2] = t.CDouble(3.0)
a.data[3] = t.CDouble(4.0)
b: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
b.data[0] = t.CDouble(5.0)
b.data[1] = t.CDouble(6.0)
b.data[2] = t.CDouble(7.0)
b.data[3] = t.CDouble(8.0)
out: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
# add
vipersimd.simd_add4d(a.data, b.data, out.data)
check("add[0]==6", out.data[0] == t.CDouble(6.0), "")
check("add[1]==8", out.data[1] == t.CDouble(8.0), "")
check("add[2]==10", out.data[2] == t.CDouble(10.0), "")
check("add[3]==12", out.data[3] == t.CDouble(12.0), "")
# sub
vipersimd.simd_sub4d(a.data, b.data, out.data)
check("sub[0]==-4", out.data[0] == t.CDouble(-4.0), "")
check("sub[3]==-4", out.data[3] == t.CDouble(-4.0), "")
# mul
vipersimd.simd_mul4d(a.data, b.data, out.data)
check("mul[0]==5", out.data[0] == t.CDouble(5.0), "")
check("mul[1]==12", out.data[1] == t.CDouble(12.0), "")
check("mul[2]==21", out.data[2] == t.CDouble(21.0), "")
check("mul[3]==32", out.data[3] == t.CDouble(32.0), "")
# div
vipersimd.simd_div4d(b.data, a.data, out.data)
check("div[0]==5", out.data[0] == t.CDouble(5.0), "")
check("div[1]==3", out.data[1] == t.CDouble(3.0), "")
a.delete()
b.delete()
out.delete()
def test_avx2_neg_abs_sqrt():
section("AVX2 neg/abs/sqrt 4d")
a: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
a.data[0] = t.CDouble(1.0)
a.data[1] = t.CDouble(-4.0)
a.data[2] = t.CDouble(9.0)
a.data[3] = t.CDouble(-16.0)
out: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
# neg
vipersimd.simd_neg4d(a.data, out.data)
check("neg[0]==-1", out.data[0] == t.CDouble(-1.0), "")
check("neg[1]==4", out.data[1] == t.CDouble(4.0), "")
check("neg[2]==-9", out.data[2] == t.CDouble(-9.0), "")
check("neg[3]==16", out.data[3] == t.CDouble(16.0), "")
# abs
vipersimd.simd_abs4d(a.data, out.data)
check("abs[0]==1", out.data[0] == t.CDouble(1.0), "")
check("abs[1]==4", out.data[1] == t.CDouble(4.0), "")
check("abs[2]==9", out.data[2] == t.CDouble(9.0), "")
check("abs[3]==16", out.data[3] == t.CDouble(16.0), "")
# sqrt
vipersimd.simd_sqrt4d(out.data, out.data)
check("sqrt[0]==1", out.data[0] == t.CDouble(1.0), "")
check("sqrt[1]==2", out.data[1] == t.CDouble(2.0), "")
check("sqrt[2]==3", out.data[2] == t.CDouble(3.0), "")
check("sqrt[3]==4", out.data[3] == t.CDouble(4.0), "")
a.delete()
out.delete()
def test_avx2_min_max():
section("AVX2 min/max 4d")
a: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
a.data[0] = t.CDouble(1.0)
a.data[1] = t.CDouble(8.0)
a.data[2] = t.CDouble(3.0)
a.data[3] = t.CDouble(6.0)
b: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
b.data[0] = t.CDouble(5.0)
b.data[1] = t.CDouble(2.0)
b.data[2] = t.CDouble(7.0)
b.data[3] = t.CDouble(4.0)
out: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
vipersimd.simd_min4d(a.data, b.data, out.data)
check("min[0]==1", out.data[0] == t.CDouble(1.0), "")
check("min[1]==2", out.data[1] == t.CDouble(2.0), "")
check("min[2]==3", out.data[2] == t.CDouble(3.0), "")
check("min[3]==4", out.data[3] == t.CDouble(4.0), "")
vipersimd.simd_max4d(a.data, b.data, out.data)
check("max[0]==5", out.data[0] == t.CDouble(5.0), "")
check("max[1]==8", out.data[1] == t.CDouble(8.0), "")
check("max[2]==7", out.data[2] == t.CDouble(7.0), "")
check("max[3]==6", out.data[3] == t.CDouble(6.0), "")
a.delete()
b.delete()
out.delete()
def test_avx2_scalar_ops():
section("AVX2 scalar broadcast ops")
a: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
a.data[0] = t.CDouble(1.0)
a.data[1] = t.CDouble(2.0)
a.data[2] = t.CDouble(3.0)
a.data[3] = t.CDouble(4.0)
s: numpy.ndarray | t.CPtr = numpy.zeros(pool, 1)
s.data[0] = t.CDouble(10.0)
out: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
vipersimd.simd_mul_scalar4d(a.data, s.data, out.data)
check("mul_scalar[0]==10", out.data[0] == t.CDouble(10.0), "")
check("mul_scalar[1]==20", out.data[1] == t.CDouble(20.0), "")
check("mul_scalar[3]==40", out.data[3] == t.CDouble(40.0), "")
vipersimd.simd_add_scalar4d(a.data, s.data, out.data)
check("add_scalar[0]==11", out.data[0] == t.CDouble(11.0), "")
check("add_scalar[3]==14", out.data[3] == t.CDouble(14.0), "")
a.delete()
s.delete()
out.delete()
def test_avx2_hsum_dot():
section("AVX2 hsum/dot 4d")
a: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
a.data[0] = t.CDouble(1.0)
a.data[1] = t.CDouble(2.0)
a.data[2] = t.CDouble(3.0)
a.data[3] = t.CDouble(4.0)
b: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
b.data[0] = t.CDouble(5.0)
b.data[1] = t.CDouble(6.0)
b.data[2] = t.CDouble(7.0)
b.data[3] = t.CDouble(8.0)
out: numpy.ndarray | t.CPtr = numpy.zeros(pool, 1)
vipersimd.simd_hsum4d(a.data, out.data)
check("hsum==10", out.data[0] == t.CDouble(10.0), "")
vipersimd.simd_dot4d(a.data, b.data, out.data)
check("dot==70", out.data[0] == t.CDouble(70.0), "")
a.delete()
b.delete()
out.delete()
def test_avx2_fma():
section("AVX2 FMA 4d")
a: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
a.data[0] = t.CDouble(1.0)
a.data[1] = t.CDouble(2.0)
a.data[2] = t.CDouble(3.0)
a.data[3] = t.CDouble(4.0)
b: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
b.data[0] = t.CDouble(5.0)
b.data[1] = t.CDouble(6.0)
b.data[2] = t.CDouble(7.0)
b.data[3] = t.CDouble(8.0)
c_val: numpy.ndarray | t.CPtr = numpy.ones(pool, 4)
out: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
vipersimd.simd_fma4d(a.data, b.data, c_val.data, out.data)
check("fma[0]==6", out.data[0] == t.CDouble(6.0), "")
check("fma[1]==13", out.data[1] == t.CDouble(13.0), "")
check("fma[2]==22", out.data[2] == t.CDouble(22.0), "")
check("fma[3]==33", out.data[3] == t.CDouble(33.0), "")
a.delete()
b.delete()
c_val.delete()
out.delete()
# ============================================================
# 2. 批量数组运算正确性测试
# ============================================================
def test_batch_array_ops():
section("Batch array operations")
n: t.CSizeT = 16
a: numpy.ndarray | t.CPtr = numpy.arange(pool, t.CDouble(1.0), t.CDouble(17.0), t.CDouble(1.0))
b: numpy.ndarray | t.CPtr = numpy.full(pool, 16, t.CDouble(2.0))
out: numpy.ndarray | t.CPtr = numpy.zeros(pool, 16)
# add array
vipersimd.simd_add_array(a.data, b.data, out.data, n)
check("add_arr[0]==3", out.data[0] == t.CDouble(3.0), "")
check("add_arr[15]==18", out.data[15] == t.CDouble(18.0), "")
# mul array
vipersimd.simd_mul_array(a.data, b.data, out.data, n)
check("mul_arr[0]==2", out.data[0] == t.CDouble(2.0), "")
check("mul_arr[3]==8", out.data[3] == t.CDouble(8.0), "")
# sqrt array
c: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
c.data[0] = t.CDouble(4.0)
c.data[1] = t.CDouble(9.0)
c.data[2] = t.CDouble(16.0)
c.data[3] = t.CDouble(25.0)
out4: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
vipersimd.simd_sqrt_array(c.data, out4.data, t.CSizeT(4))
check("sqrt_arr[0]==2", out4.data[0] == t.CDouble(2.0), "")
check("sqrt_arr[3]==5", out4.data[3] == t.CDouble(5.0), "")
# abs array
d: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
d.data[0] = t.CDouble(-1.0)
d.data[1] = t.CDouble(-2.0)
d.data[2] = t.CDouble(3.0)
d.data[3] = t.CDouble(-4.0)
vipersimd.simd_abs_array(d.data, out4.data, t.CSizeT(4))
check("abs_arr[0]==1", out4.data[0] == t.CDouble(1.0), "")
check("abs_arr[1]==2", out4.data[1] == t.CDouble(2.0), "")
check("abs_arr[3]==4", out4.data[3] == t.CDouble(4.0), "")
# dot array
e: numpy.ndarray | t.CPtr = numpy.zeros(pool, 4)
e.data[0] = t.CDouble(1.0)
e.data[1] = t.CDouble(2.0)
e.data[2] = t.CDouble(3.0)
e.data[3] = t.CDouble(4.0)
f: numpy.ndarray | t.CPtr = numpy.ones(pool, 4)
out1: numpy.ndarray | t.CPtr = numpy.zeros(pool, 1)
vipersimd.simd_dot_array(e.data, f.data, t.CSizeT(4), out1.data)
check("dot_arr==10", out1.data[0] == t.CDouble(10.0), "")
a.delete()
b.delete()
out.delete()
c.delete()
out4.delete()
d.delete()
e.delete()
f.delete()
out1.delete()
# ============================================================
# 3. LLVMIR 精确标量运算测试
# ============================================================
def test_llvmir_scalar():
section("LLVMIR precise scalar ops")
# sqrt
r: t.CDouble = vipersimd.simd_sqrt(t.CDouble(4.0))
check("sqrt(4)==2", r == t.CDouble(2.0), "")
# fabs
r = vipersimd.simd_fabs(t.CDouble(-3.5))
check("fabs(-3.5)==3.5", r == t.CDouble(3.5), "")
# floor
r = vipersimd.simd_floor(t.CDouble(3.7))
check("floor(3.7)==3", r == t.CDouble(3.0), "")
# ceil
r = vipersimd.simd_ceil(t.CDouble(3.2))
check("ceil(3.2)==4", r == t.CDouble(4.0), "")
# round
r = vipersimd.simd_round(t.CDouble(3.5))
check("round(3.5)==4", r == t.CDouble(4.0), "")
# trunc
r = vipersimd.simd_trunc(t.CDouble(3.9))
check("trunc(3.9)==3", r == t.CDouble(3.0), "")
# fma: 2*3+1 = 7
r = vipersimd.simd_fma(t.CDouble(2.0), t.CDouble(3.0), t.CDouble(1.0))
check("fma(2,3,1)==7", r == t.CDouble(7.0), "")
# copysign
r = vipersimd.simd_copysign(t.CDouble(3.0), t.CDouble(-1.0))
check("copysign(3,-1)==-3", r == t.CDouble(-3.0), "")
# neg
r = vipersimd.simd_neg(t.CDouble(1.5))
check("neg(1.5)==-1.5", r == t.CDouble(-1.5), "")
# minnum
r = vipersimd.simd_minnum(t.CDouble(2.0), t.CDouble(3.0))
check("minnum(2,3)==2", r == t.CDouble(2.0), "")
# maxnum
r = vipersimd.simd_maxnum(t.CDouble(2.0), t.CDouble(3.0))
check("maxnum(2,3)==3", r == t.CDouble(3.0), "")
# exp
r = vipersimd.simd_exp(t.CDouble(1.0))
check("exp(1)~=2.718", r > t.CDouble(2.71) and r < t.CDouble(2.72), "")
# log
r = vipersimd.simd_log(t.CDouble(2.718281828459045))
check("log(e)~=1", r > t.CDouble(0.999) and r < t.CDouble(1.001), "")
# sin
r = vipersimd.simd_sin(t.CDouble(0.0))
check("sin(0)~=0", r > t.CDouble(-0.001) and r < t.CDouble(0.001), "")
# cos
r = vipersimd.simd_cos(t.CDouble(0.0))
check("cos(0)~=1", r > t.CDouble(0.999) and r < t.CDouble(1.001), "")
# pow
r = vipersimd.simd_pow(t.CDouble(2.0), t.CDouble(10.0))
check("pow(2,10)==1024", r == t.CDouble(1024.0), "")
# ============================================================
# 4. 性能基准测试: SIMD vs 标量
# ============================================================
N_ELEM: t.CDefine = 40000
N_ITER: t.CDefine = 50000
def scalar_add_array(a: t.CPtr, b: t.CPtr, out: t.CPtr, n: t.CSizeT):
i: t.CSizeT = 0
da: t.CDouble | t.CPtr = a
db: t.CDouble | t.CPtr = b
dout: t.CDouble | t.CPtr = out
while i < n:
dout[i] = da[i] + db[i]
i += 1
def scalar_mul_array(a: t.CPtr, b: t.CPtr, out: t.CPtr, n: t.CSizeT):
i: t.CSizeT = 0
da: t.CDouble | t.CPtr = a
db: t.CDouble | t.CPtr = b
dout: t.CDouble | t.CPtr = out
while i < n:
dout[i] = da[i] * db[i]
i += 1
def scalar_sqrt_array(a: t.CPtr, out: t.CPtr, n: t.CSizeT):
i: t.CSizeT = 0
da: t.CDouble | t.CPtr = a
dout: t.CDouble | t.CPtr = out
while i < n:
dout[i] = vipersimd.simd_sqrt(da[i])
i += 1
def benchmark_simd_vs_scalar():
section("Benchmark: SIMD vs Scalar")
a: numpy.ndarray | t.CPtr = numpy.full(pool, N_ELEM, t.CDouble(1.5))
b: numpy.ndarray | t.CPtr = numpy.full(pool, N_ELEM, t.CDouble(2.5))
out: numpy.ndarray | t.CPtr = numpy.zeros(pool, N_ELEM)
# --- SIMD add benchmark ---
t0: t.CLong = clock()
iter: t.CInt = 0
while iter < N_ITER:
vipersimd.simd_add_array(a.data, b.data, out.data, t.CSizeT(N_ELEM))
iter += 1
t1: t.CLong = clock()
simd_add_ms: t.CLong = t1 - t0
# --- Scalar add benchmark ---
t0 = clock()
iter = 0
while iter < N_ITER:
scalar_add_array(a.data, b.data, out.data, t.CSizeT(N_ELEM))
iter += 1
t1 = clock()
scalar_add_ms: t.CLong = t1 - t0
printf(" ADD: SIMD=%ldms Scalar=%ldms", simd_add_ms, scalar_add_ms)
if scalar_add_ms > 0:
printf(" speedup=%.1fx\n", t.CDouble(scalar_add_ms) / t.CDouble(simd_add_ms))
else:
printf(" speedup=N/A\n")
# --- SIMD mul benchmark ---
t0 = clock()
iter = 0
while iter < N_ITER:
vipersimd.simd_mul_array(a.data, b.data, out.data, t.CSizeT(N_ELEM))
iter += 1
t1 = clock()
simd_mul_ms: t.CLong = t1 - t0
# --- Scalar mul benchmark ---
t0 = clock()
iter = 0
while iter < N_ITER:
scalar_mul_array(a.data, b.data, out.data, t.CSizeT(N_ELEM))
iter += 1
t1 = clock()
scalar_mul_ms: t.CLong = t1 - t0
printf(" MUL: SIMD=%ldms Scalar=%ldms", simd_mul_ms, scalar_mul_ms)
if scalar_mul_ms > 0:
printf(" speedup=%.1fx\n", t.CDouble(scalar_mul_ms) / t.CDouble(simd_mul_ms))
else:
printf(" speedup=N/A\n")
# --- SIMD sqrt benchmark ---
t0 = clock()
iter = 0
while iter < N_ITER:
vipersimd.simd_sqrt_array(a.data, out.data, t.CSizeT(N_ELEM))
iter += 1
t1 = clock()
simd_sqrt_ms: t.CLong = t1 - t0
# --- Scalar sqrt benchmark ---
t0 = clock()
iter = 0
while iter < N_ITER:
scalar_sqrt_array(a.data, out.data, t.CSizeT(N_ELEM))
iter += 1
t1 = clock()
scalar_sqrt_ms: t.CLong = t1 - t0
printf(" SQRT: SIMD=%ldms Scalar=%ldms", simd_sqrt_ms, scalar_sqrt_ms)
if scalar_sqrt_ms > 0:
printf(" speedup=%.1fx\n", t.CDouble(scalar_sqrt_ms) / t.CDouble(simd_sqrt_ms))
else:
printf(" speedup=N/A\n")
# 验证结果正确性
check("bench result[0] valid", out.data[0] > t.CDouble(0.0), "")
a.delete()
b.delete()
out.delete()
# ============================================================
# Main
# ============================================================
def main() -> t.CInt:
global arena, pool
arena = calloc(4096 * 1024, 1) # 4MB arena
pool = mpool.MPool(arena, 4096 * 1024, 0) # bump allocator
printf("=== ViperSIMD Test Suite ===\n")
test_avx2_add_sub_mul_div()
test_avx2_neg_abs_sqrt()
test_avx2_min_max()
test_avx2_scalar_ops()
test_avx2_hsum_dot()
test_avx2_fma()
test_batch_array_ops()
test_llvmir_scalar()
benchmark_simd_vs_scalar()
printf("\n=== Results: %d passed, %d failed ===\n", test_passed, test_failed)
return 0