793 lines
24 KiB
Plaintext
793 lines
24 KiB
Plaintext
import t, c
|
||
|
||
|
||
# ============================================================
|
||
# AVX2 double precision vector operations (4 x double, 256-bit YMM)
|
||
# ============================================================
|
||
|
||
def simd_add4d(a: t.CPtr, b: t.CPtr, out: t.CPtr):
|
||
"""AVX2: out[0:4] = a[0:4] + b[0:4]"""
|
||
c.Asm(f"""vmovupd ymm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
vmovupd ymm1, [{c.AsmInp(b, t.ASM_DESCR.REG_ANY)}]
|
||
vaddpd ymm0, ymm0, ymm1
|
||
vmovupd [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], ymm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1'])
|
||
|
||
|
||
def simd_sub4d(a: t.CPtr, b: t.CPtr, out: t.CPtr):
|
||
"""AVX2: out[0:4] = a[0:4] - b[0:4]"""
|
||
c.Asm(f"""vmovupd ymm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
vmovupd ymm1, [{c.AsmInp(b, t.ASM_DESCR.REG_ANY)}]
|
||
vsubpd ymm0, ymm0, ymm1
|
||
vmovupd [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], ymm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1'])
|
||
|
||
|
||
def simd_mul4d(a: t.CPtr, b: t.CPtr, out: t.CPtr):
|
||
"""AVX2: out[0:4] = a[0:4] * b[0:4]"""
|
||
c.Asm(f"""vmovupd ymm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
vmovupd ymm1, [{c.AsmInp(b, t.ASM_DESCR.REG_ANY)}]
|
||
vmulpd ymm0, ymm0, ymm1
|
||
vmovupd [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], ymm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1'])
|
||
|
||
|
||
def simd_div4d(a: t.CPtr, b: t.CPtr, out: t.CPtr):
|
||
"""AVX2: out[0:4] = a[0:4] / b[0:4]"""
|
||
c.Asm(f"""vmovupd ymm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
vmovupd ymm1, [{c.AsmInp(b, t.ASM_DESCR.REG_ANY)}]
|
||
vdivpd ymm0, ymm0, ymm1
|
||
vmovupd [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], ymm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1'])
|
||
|
||
|
||
def simd_sqrt4d(a: t.CPtr, out: t.CPtr):
|
||
"""AVX2: out[0:4] = sqrt(a[0:4])"""
|
||
c.Asm(f"""vmovupd ymm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
vsqrtpd ymm0, ymm0
|
||
vmovupd [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], ymm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0'])
|
||
|
||
|
||
def simd_neg4d(a: t.CPtr, out: t.CPtr):
|
||
"""AVX2: out[0:4] = -a[0:4] (flip sign bit via XOR)"""
|
||
c.Asm(f"""vmovupd ymm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
vpcmpeqd ymm1, ymm1, ymm1
|
||
vpsllq ymm1, ymm1, 63
|
||
vxorpd ymm0, ymm0, ymm1
|
||
vmovupd [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], ymm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1'])
|
||
|
||
|
||
def simd_abs4d(a: t.CPtr, out: t.CPtr):
|
||
"""AVX2: out[0:4] = |a[0:4]| (clear sign bit via ANDN)"""
|
||
c.Asm(f"""vmovupd ymm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
vpcmpeqd ymm1, ymm1, ymm1
|
||
vpsllq ymm1, ymm1, 63
|
||
vandnpd ymm0, ymm1, ymm0
|
||
vmovupd [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], ymm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1'])
|
||
|
||
|
||
def simd_max4d(a: t.CPtr, b: t.CPtr, out: t.CPtr):
|
||
"""AVX2: out[0:4] = max(a[0:4], b[0:4])"""
|
||
c.Asm(f"""vmovupd ymm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
vmovupd ymm1, [{c.AsmInp(b, t.ASM_DESCR.REG_ANY)}]
|
||
vmaxpd ymm0, ymm0, ymm1
|
||
vmovupd [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], ymm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1'])
|
||
|
||
|
||
def simd_min4d(a: t.CPtr, b: t.CPtr, out: t.CPtr):
|
||
"""AVX2: out[0:4] = min(a[0:4], b[0:4])"""
|
||
c.Asm(f"""vmovupd ymm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
vmovupd ymm1, [{c.AsmInp(b, t.ASM_DESCR.REG_ANY)}]
|
||
vminpd ymm0, ymm0, ymm1
|
||
vmovupd [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], ymm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1'])
|
||
|
||
|
||
def simd_mul_scalar4d(a: t.CPtr, scalar: t.CPtr, out: t.CPtr):
|
||
"""AVX2: out[0:4] = a[0:4] * scalar (broadcast)"""
|
||
c.Asm(f"""vmovupd ymm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
vbroadcastsd ymm1, [{c.AsmInp(scalar, t.ASM_DESCR.REG_ANY)}]
|
||
vmulpd ymm0, ymm0, ymm1
|
||
vmovupd [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], ymm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1'])
|
||
|
||
|
||
def simd_add_scalar4d(a: t.CPtr, scalar: t.CPtr, out: t.CPtr):
|
||
"""AVX2: out[0:4] = a[0:4] + scalar (broadcast)"""
|
||
c.Asm(f"""vmovupd ymm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
vbroadcastsd ymm1, [{c.AsmInp(scalar, t.ASM_DESCR.REG_ANY)}]
|
||
vaddpd ymm0, ymm0, ymm1
|
||
vmovupd [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], ymm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1'])
|
||
|
||
|
||
def simd_hsum4d(a: t.CPtr, out: t.CPtr):
|
||
"""AVX2: out[0] = a[0]+a[1]+a[2]+a[3] (horizontal sum)"""
|
||
c.Asm(f"""vmovupd ymm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
vextractf128 xmm1, ymm0, 1
|
||
vaddpd xmm0, xmm0, xmm1
|
||
vhaddpd xmm0, xmm0, xmm0
|
||
vmovsd [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], xmm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1'])
|
||
|
||
|
||
def simd_dot4d(a: t.CPtr, b: t.CPtr, out: t.CPtr):
|
||
"""AVX2: out[0] = dot(a[0:4], b[0:4])"""
|
||
c.Asm(f"""vmovupd ymm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
vmovupd ymm1, [{c.AsmInp(b, t.ASM_DESCR.REG_ANY)}]
|
||
vmulpd ymm0, ymm0, ymm1
|
||
vextractf128 xmm1, ymm0, 1
|
||
vaddpd xmm0, xmm0, xmm1
|
||
vhaddpd xmm0, xmm0, xmm0
|
||
vmovsd [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], xmm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1'])
|
||
|
||
|
||
def simd_fma4d(a: t.CPtr, b: t.CPtr, c_val: t.CPtr, out: t.CPtr):
|
||
"""AVX2: out[0:4] = a[0:4]*b[0:4] + c[0:4] (fused multiply-add)
|
||
Uses vfmadd213pd: ymm0 = ymm1 * ymm0 + ymm2
|
||
Load order: b->ymm0, a->ymm1, c->ymm2"""
|
||
c.Asm(f"""vmovupd ymm0, [{c.AsmInp(b, t.ASM_DESCR.REG_ANY)}]
|
||
vmovupd ymm1, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
vmovupd ymm2, [{c.AsmInp(c_val, t.ASM_DESCR.REG_ANY)}]
|
||
vfmadd213pd ymm0, ymm1, ymm2
|
||
vmovupd [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], ymm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1', 'ymm2'])
|
||
|
||
|
||
# ============================================================
|
||
# SSE2 single precision vector operations (4 x float, 128-bit XMM)
|
||
# ============================================================
|
||
|
||
def simd_add4f(a: t.CPtr, b: t.CPtr, out: t.CPtr):
|
||
"""SSE2: out[0:4] = a[0:4] + b[0:4] (float)"""
|
||
c.Asm(f"""movups xmm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
movups xmm1, [{c.AsmInp(b, t.ASM_DESCR.REG_ANY)}]
|
||
addps xmm0, xmm1
|
||
movups [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], xmm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'xmm0', 'xmm1'])
|
||
|
||
|
||
def simd_sub4f(a: t.CPtr, b: t.CPtr, out: t.CPtr):
|
||
"""SSE2: out[0:4] = a[0:4] - b[0:4] (float)"""
|
||
c.Asm(f"""movups xmm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
movups xmm1, [{c.AsmInp(b, t.ASM_DESCR.REG_ANY)}]
|
||
subps xmm0, xmm1
|
||
movups [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], xmm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'xmm0', 'xmm1'])
|
||
|
||
|
||
def simd_mul4f(a: t.CPtr, b: t.CPtr, out: t.CPtr):
|
||
"""SSE2: out[0:4] = a[0:4] * b[0:4] (float)"""
|
||
c.Asm(f"""movups xmm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
movups xmm1, [{c.AsmInp(b, t.ASM_DESCR.REG_ANY)}]
|
||
mulps xmm0, xmm1
|
||
movups [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], xmm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'xmm0', 'xmm1'])
|
||
|
||
|
||
def simd_div4f(a: t.CPtr, b: t.CPtr, out: t.CPtr):
|
||
"""SSE2: out[0:4] = a[0:4] / b[0:4] (float)"""
|
||
c.Asm(f"""movups xmm0, [{c.AsmInp(a, t.ASM_DESCR.REG_ANY)}]
|
||
movups xmm1, [{c.AsmInp(b, t.ASM_DESCR.REG_ANY)}]
|
||
divps xmm0, xmm1
|
||
movups [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], xmm0""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'xmm0', 'xmm1'])
|
||
|
||
|
||
# ============================================================
|
||
# Batch array operations (AVX2 loop, process N double values)
|
||
# ============================================================
|
||
|
||
def simd_add_array(a: t.CPtr, b: t.CPtr, out: t.CPtr, n: t.CSizeT):
|
||
"""AVX2 batch: out[i:i+4] = a[i:i+4] + b[i:i+4], i=0,4,8,...
|
||
×8 双缓冲（32 doubles/大迭代）：a(ymm0-7)+b(ymm8-15) 16 路 load 全部在飞后
|
||
做 8 条独立 add→store 链，最大化内存级并行、压榨访存口。
|
||
尾巴（不足 8 个 ymm 的块）回退单向量循环，语义与原版一致（floor(n/4) 向量）。
|
||
注：ymm0-15 clobber 需 TPV c.Asm 支持（CallGenCBuiltin ASM_MAX_CLOBBERS=32）。"""
|
||
c.Asm(f"""mov rcx, {c.AsmInp(n, t.ASM_DESCR.REG_ANY)}
|
||
shr rcx, 2
|
||
mov r8, {c.AsmInp(a, t.ASM_DESCR.REG_ANY)}
|
||
mov r9, {c.AsmInp(b, t.ASM_DESCR.REG_ANY)}
|
||
mov r10, {c.AsmInp(out, t.ASM_DESCR.REG_ANY)}
|
||
mov r11, rcx
|
||
shr r11, 3
|
||
test r11, r11
|
||
jz .La_add_tail
|
||
.La_add_big:
|
||
vmovupd ymm0, [r8]
|
||
vmovupd ymm1, [r8 + 32]
|
||
vmovupd ymm2, [r8 + 64]
|
||
vmovupd ymm3, [r8 + 96]
|
||
vmovupd ymm4, [r8 + 128]
|
||
vmovupd ymm5, [r8 + 160]
|
||
vmovupd ymm6, [r8 + 192]
|
||
vmovupd ymm7, [r8 + 224]
|
||
vmovupd ymm8, [r9]
|
||
vmovupd ymm9, [r9 + 32]
|
||
vmovupd ymm10, [r9 + 64]
|
||
vmovupd ymm11, [r9 + 96]
|
||
vmovupd ymm12, [r9 + 128]
|
||
vmovupd ymm13, [r9 + 160]
|
||
vmovupd ymm14, [r9 + 192]
|
||
vmovupd ymm15, [r9 + 224]
|
||
vaddpd ymm0, ymm0, ymm8
|
||
vaddpd ymm1, ymm1, ymm9
|
||
vaddpd ymm2, ymm2, ymm10
|
||
vaddpd ymm3, ymm3, ymm11
|
||
vaddpd ymm4, ymm4, ymm12
|
||
vaddpd ymm5, ymm5, ymm13
|
||
vaddpd ymm6, ymm6, ymm14
|
||
vaddpd ymm7, ymm7, ymm15
|
||
vmovupd [r10], ymm0
|
||
vmovupd [r10 + 32], ymm1
|
||
vmovupd [r10 + 64], ymm2
|
||
vmovupd [r10 + 96], ymm3
|
||
vmovupd [r10 + 128], ymm4
|
||
vmovupd [r10 + 160], ymm5
|
||
vmovupd [r10 + 192], ymm6
|
||
vmovupd [r10 + 224], ymm7
|
||
add r8, 256
|
||
add r9, 256
|
||
add r10, 256
|
||
dec r11
|
||
jnz .La_add_big
|
||
.La_add_tail:
|
||
and rcx, 7
|
||
jz .La_add_end
|
||
.La_add_tloop:
|
||
vmovupd ymm0, [r8]
|
||
vmovupd ymm1, [r9]
|
||
vaddpd ymm0, ymm0, ymm1
|
||
vmovupd [r10], ymm0
|
||
add r8, 32
|
||
add r9, 32
|
||
add r10, 32
|
||
dec rcx
|
||
jnz .La_add_tloop
|
||
.La_add_end:
|
||
vzeroupper""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1', 'ymm2', 'ymm3',
|
||
'ymm4', 'ymm5', 'ymm6', 'ymm7', 'ymm8', 'ymm9', 'ymm10', 'ymm11',
|
||
'ymm12', 'ymm13', 'ymm14', 'ymm15', 'rcx', 'r8', 'r9', 'r10', 'r11'])
|
||
|
||
|
||
def simd_sub_array(a: t.CPtr, b: t.CPtr, out: t.CPtr, n: t.CSizeT):
|
||
"""AVX2 batch: out[i:i+4] = a[i:i+4] - b[i:i+4]"""
|
||
c.Asm(f"""mov rcx, {c.AsmInp(n, t.ASM_DESCR.REG_ANY)}
|
||
shr rcx, 2
|
||
mov r8, {c.AsmInp(a, t.ASM_DESCR.REG_ANY)}
|
||
mov r9, {c.AsmInp(b, t.ASM_DESCR.REG_ANY)}
|
||
mov r10, {c.AsmInp(out, t.ASM_DESCR.REG_ANY)}
|
||
test rcx, rcx
|
||
jz .Lsub_end
|
||
.Lsub_loop:
|
||
vmovupd ymm0, [r8]
|
||
vmovupd ymm1, [r9]
|
||
vsubpd ymm0, ymm0, ymm1
|
||
vmovupd [r10], ymm0
|
||
add r8, 32
|
||
add r9, 32
|
||
add r10, 32
|
||
dec rcx
|
||
jnz .Lsub_loop
|
||
.Lsub_end:
|
||
vzeroupper""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1', 'rcx', 'r8', 'r9', 'r10'])
|
||
|
||
|
||
def simd_mul_array(a: t.CPtr, b: t.CPtr, out: t.CPtr, n: t.CSizeT):
|
||
"""AVX2 batch: out[i:i+4] = a[i:i+4] * b[i:i+4]
|
||
×8 双缓冲（同 simd_add_array）。"""
|
||
c.Asm(f"""mov rcx, {c.AsmInp(n, t.ASM_DESCR.REG_ANY)}
|
||
shr rcx, 2
|
||
mov r8, {c.AsmInp(a, t.ASM_DESCR.REG_ANY)}
|
||
mov r9, {c.AsmInp(b, t.ASM_DESCR.REG_ANY)}
|
||
mov r10, {c.AsmInp(out, t.ASM_DESCR.REG_ANY)}
|
||
mov r11, rcx
|
||
shr r11, 3
|
||
test r11, r11
|
||
jz .Lm_mul_tail
|
||
.Lm_mul_big:
|
||
vmovupd ymm0, [r8]
|
||
vmovupd ymm1, [r8 + 32]
|
||
vmovupd ymm2, [r8 + 64]
|
||
vmovupd ymm3, [r8 + 96]
|
||
vmovupd ymm4, [r8 + 128]
|
||
vmovupd ymm5, [r8 + 160]
|
||
vmovupd ymm6, [r8 + 192]
|
||
vmovupd ymm7, [r8 + 224]
|
||
vmovupd ymm8, [r9]
|
||
vmovupd ymm9, [r9 + 32]
|
||
vmovupd ymm10, [r9 + 64]
|
||
vmovupd ymm11, [r9 + 96]
|
||
vmovupd ymm12, [r9 + 128]
|
||
vmovupd ymm13, [r9 + 160]
|
||
vmovupd ymm14, [r9 + 192]
|
||
vmovupd ymm15, [r9 + 224]
|
||
vmulpd ymm0, ymm0, ymm8
|
||
vmulpd ymm1, ymm1, ymm9
|
||
vmulpd ymm2, ymm2, ymm10
|
||
vmulpd ymm3, ymm3, ymm11
|
||
vmulpd ymm4, ymm4, ymm12
|
||
vmulpd ymm5, ymm5, ymm13
|
||
vmulpd ymm6, ymm6, ymm14
|
||
vmulpd ymm7, ymm7, ymm15
|
||
vmovupd [r10], ymm0
|
||
vmovupd [r10 + 32], ymm1
|
||
vmovupd [r10 + 64], ymm2
|
||
vmovupd [r10 + 96], ymm3
|
||
vmovupd [r10 + 128], ymm4
|
||
vmovupd [r10 + 160], ymm5
|
||
vmovupd [r10 + 192], ymm6
|
||
vmovupd [r10 + 224], ymm7
|
||
add r8, 256
|
||
add r9, 256
|
||
add r10, 256
|
||
dec r11
|
||
jnz .Lm_mul_big
|
||
.Lm_mul_tail:
|
||
and rcx, 7
|
||
jz .Lm_mul_end
|
||
.Lm_mul_tloop:
|
||
vmovupd ymm0, [r8]
|
||
vmovupd ymm1, [r9]
|
||
vmulpd ymm0, ymm0, ymm1
|
||
vmovupd [r10], ymm0
|
||
add r8, 32
|
||
add r9, 32
|
||
add r10, 32
|
||
dec rcx
|
||
jnz .Lm_mul_tloop
|
||
.Lm_mul_end:
|
||
vzeroupper""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1', 'ymm2', 'ymm3',
|
||
'ymm4', 'ymm5', 'ymm6', 'ymm7', 'ymm8', 'ymm9', 'ymm10', 'ymm11',
|
||
'ymm12', 'ymm13', 'ymm14', 'ymm15', 'rcx', 'r8', 'r9', 'r10', 'r11'])
|
||
|
||
|
||
def simd_fma_array(a: t.CPtr, b: t.CPtr, c_val: t.CPtr, out: t.CPtr, n: t.CSizeT):
|
||
"""AVX2 batch (单趟融合乘加): out[i:i+4] = a[i:i+4]*b[i:i+4] + c_val[i:i+4]
|
||
×8 双缓冲（同 simd_mul_array）：a(ymm0-7)+b(ymm8-15) 16 路 load 在飞后，
|
||
用 8 条 vfmadd213pd 以 c_val 内存操作数为加数一次完成 FMA 并落 out，
|
||
避免 mul→tmp→add 的两遍写法（读流量 = 每元素 24B，写 8B，无中间数组）。
|
||
语义：vfmadd213pd ymmX, ymmY, mem => ymmX = ymmY*ymmX + mem
|
||
ymm8..15 = a(big lane)*b(big lane) + c(big lane)，单次舍入。
|
||
尾巴（不足 8 个 ymm 的块）回退单向量循环。n 须为 4 的倍数（floor(n/4) 向量）。"""
|
||
c.Asm(f"""mov rcx, {c.AsmInp(n, t.ASM_DESCR.REG_ANY)}
|
||
shr rcx, 2
|
||
mov r8, {c.AsmInp(a, t.ASM_DESCR.REG_ANY)}
|
||
mov r9, {c.AsmInp(b, t.ASM_DESCR.REG_ANY)}
|
||
mov rdx, {c.AsmInp(c_val, t.ASM_DESCR.REG_ANY)}
|
||
mov r10, {c.AsmInp(out, t.ASM_DESCR.REG_ANY)}
|
||
mov r11, rcx
|
||
shr r11, 3
|
||
test r11, r11
|
||
jz .Lfma_tail
|
||
.Lfma_big:
|
||
vmovupd ymm0, [r8]
|
||
vmovupd ymm1, [r8 + 32]
|
||
vmovupd ymm2, [r8 + 64]
|
||
vmovupd ymm3, [r8 + 96]
|
||
vmovupd ymm4, [r8 + 128]
|
||
vmovupd ymm5, [r8 + 160]
|
||
vmovupd ymm6, [r8 + 192]
|
||
vmovupd ymm7, [r8 + 224]
|
||
vmovupd ymm8, [r9]
|
||
vmovupd ymm9, [r9 + 32]
|
||
vmovupd ymm10, [r9 + 64]
|
||
vmovupd ymm11, [r9 + 96]
|
||
vmovupd ymm12, [r9 + 128]
|
||
vmovupd ymm13, [r9 + 160]
|
||
vmovupd ymm14, [r9 + 192]
|
||
vmovupd ymm15, [r9 + 224]
|
||
vfmadd213pd ymm8, ymm0, [rdx]
|
||
vfmadd213pd ymm9, ymm1, [rdx + 32]
|
||
vfmadd213pd ymm10, ymm2, [rdx + 64]
|
||
vfmadd213pd ymm11, ymm3, [rdx + 96]
|
||
vfmadd213pd ymm12, ymm4, [rdx + 128]
|
||
vfmadd213pd ymm13, ymm5, [rdx + 160]
|
||
vfmadd213pd ymm14, ymm6, [rdx + 192]
|
||
vfmadd213pd ymm15, ymm7, [rdx + 224]
|
||
vmovupd [r10], ymm8
|
||
vmovupd [r10 + 32], ymm9
|
||
vmovupd [r10 + 64], ymm10
|
||
vmovupd [r10 + 96], ymm11
|
||
vmovupd [r10 + 128], ymm12
|
||
vmovupd [r10 + 160], ymm13
|
||
vmovupd [r10 + 192], ymm14
|
||
vmovupd [r10 + 224], ymm15
|
||
add r8, 256
|
||
add r9, 256
|
||
add rdx, 256
|
||
add r10, 256
|
||
dec r11
|
||
jnz .Lfma_big
|
||
.Lfma_tail:
|
||
and rcx, 7
|
||
jz .Lfma_end
|
||
.Lfma_tloop:
|
||
vmovupd ymm0, [r8]
|
||
vmovupd ymm1, [r9]
|
||
vfmadd213pd ymm0, ymm1, [rdx]
|
||
vmovupd [r10], ymm0
|
||
add r8, 32
|
||
add r9, 32
|
||
add rdx, 32
|
||
add r10, 32
|
||
dec rcx
|
||
jnz .Lfma_tloop
|
||
.Lfma_end:
|
||
vzeroupper""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1', 'ymm2', 'ymm3',
|
||
'ymm4', 'ymm5', 'ymm6', 'ymm7', 'ymm8', 'ymm9', 'ymm10', 'ymm11',
|
||
'ymm12', 'ymm13', 'ymm14', 'ymm15', 'rcx', 'rdx', 'r8', 'r9', 'r10', 'r11'])
|
||
|
||
|
||
def scalar_fma_array(a: t.CPtr, b: t.CPtr, c_val: t.CPtr, out: t.CPtr, n: t.CSizeT):
|
||
"""标量(逐元素)融合乘加: out[i] = a[i]*b[i] + c_val[i], 一次只算 1 个 double。
|
||
|
||
与 simd_fma_array 是同一数学内核的"未向量化"对照实现:
|
||
* 指令用标量 FMA `vfmadd213sd`(xmm 低位), 与 SIMD 版的 `vfmadd213pd`
|
||
同为硬件融合乘加、同为单次舍入 -> 逐元素结果与 SIMD 版**逐位一致**,
|
||
CHECK 可直接对照;
|
||
* 每次循环 3 读 1 写、8 字节粒度, 相比 ×4 的 ymm 少 4 倍数据并行度,
|
||
用来量化"向量化本身"带来的收益(SIMD 加速比)。
|
||
n 须 >= 1; n == 0 直接跳过。"""
|
||
c.Asm(f"""mov rcx, {c.AsmInp(n, t.ASM_DESCR.REG_ANY)}
|
||
mov r8, {c.AsmInp(a, t.ASM_DESCR.REG_ANY)}
|
||
mov r9, {c.AsmInp(b, t.ASM_DESCR.REG_ANY)}
|
||
mov rdx, {c.AsmInp(c_val, t.ASM_DESCR.REG_ANY)}
|
||
mov r10, {c.AsmInp(out, t.ASM_DESCR.REG_ANY)}
|
||
test rcx, rcx
|
||
jz .Lsca_end
|
||
.Lsca_loop:
|
||
vmovsd xmm0, [r8]
|
||
vmovsd xmm1, [r9]
|
||
vfmadd213sd xmm0, xmm1, [rdx]
|
||
vmovsd [r10], xmm0
|
||
add r8, 8
|
||
add r9, 8
|
||
add rdx, 8
|
||
add r10, 8
|
||
dec rcx
|
||
jnz .Lsca_loop
|
||
.Lsca_end:""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'xmm0', 'xmm1',
|
||
'rcx', 'rdx', 'r8', 'r9', 'r10'])
|
||
|
||
|
||
def scalar_add_array(a: t.CPtr, b: t.CPtr, out: t.CPtr, n: t.CSizeT):
|
||
"""标量(逐元素)加法: out[i] = a[i] + b[i], 一次 1 个 double。
|
||
|
||
`simd_add_array` 的未向量化对照: 标量 `vaddsd` vs SIMD `vaddpd`,
|
||
同为 IEEE 加法 -> 结果逐位一致。"""
|
||
c.Asm(f"""mov rcx, {c.AsmInp(n, t.ASM_DESCR.REG_ANY)}
|
||
mov r8, {c.AsmInp(a, t.ASM_DESCR.REG_ANY)}
|
||
mov r9, {c.AsmInp(b, t.ASM_DESCR.REG_ANY)}
|
||
mov r10, {c.AsmInp(out, t.ASM_DESCR.REG_ANY)}
|
||
test rcx, rcx
|
||
jz .Lsca_add_end
|
||
.Lsca_add_loop:
|
||
vmovsd xmm0, [r8]
|
||
vaddsd xmm0, xmm0, [r9]
|
||
vmovsd [r10], xmm0
|
||
add r8, 8
|
||
add r9, 8
|
||
add r10, 8
|
||
dec rcx
|
||
jnz .Lsca_add_loop
|
||
.Lsca_add_end:""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'xmm0', 'rcx', 'r8', 'r9', 'r10'])
|
||
|
||
|
||
def scalar_mul_array(a: t.CPtr, b: t.CPtr, out: t.CPtr, n: t.CSizeT):
|
||
"""标量(逐元素)乘法: out[i] = a[i] * b[i], 一次 1 个 double。
|
||
|
||
`simd_mul_array` 的未向量化对照: 标量 `vmulsd` vs SIMD `vmulpd`,
|
||
同为 IEEE 乘法 -> 结果逐位一致。"""
|
||
c.Asm(f"""mov rcx, {c.AsmInp(n, t.ASM_DESCR.REG_ANY)}
|
||
mov r8, {c.AsmInp(a, t.ASM_DESCR.REG_ANY)}
|
||
mov r9, {c.AsmInp(b, t.ASM_DESCR.REG_ANY)}
|
||
mov r10, {c.AsmInp(out, t.ASM_DESCR.REG_ANY)}
|
||
test rcx, rcx
|
||
jz .Lsca_mul_end
|
||
.Lsca_mul_loop:
|
||
vmovsd xmm0, [r8]
|
||
vmulsd xmm0, xmm0, [r9]
|
||
vmovsd [r10], xmm0
|
||
add r8, 8
|
||
add r9, 8
|
||
add r10, 8
|
||
dec rcx
|
||
jnz .Lsca_mul_loop
|
||
.Lsca_mul_end:""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'xmm0', 'rcx', 'r8', 'r9', 'r10'])
|
||
|
||
|
||
def scalar_sqrt_array(a: t.CPtr, out: t.CPtr, n: t.CSizeT):
|
||
"""标量(逐元素)平方根: out[i] = sqrt(a[i]), 一次 1 个 double。
|
||
|
||
`simd_sqrt_array` 的未向量化对照: 标量 `vsqrtsd` vs SIMD `vsqrtpd`,
|
||
两者都是 IEEE 正确舍入的平方根 -> 结果逐位一致。"""
|
||
c.Asm(f"""mov rcx, {c.AsmInp(n, t.ASM_DESCR.REG_ANY)}
|
||
mov r8, {c.AsmInp(a, t.ASM_DESCR.REG_ANY)}
|
||
mov r10, {c.AsmInp(out, t.ASM_DESCR.REG_ANY)}
|
||
test rcx, rcx
|
||
jz .Lsca_sqrt_end
|
||
.Lsca_sqrt_loop:
|
||
vmovsd xmm0, [r8]
|
||
vsqrtsd xmm0, xmm0, xmm0
|
||
vmovsd [r10], xmm0
|
||
add r8, 8
|
||
add r10, 8
|
||
dec rcx
|
||
jnz .Lsca_sqrt_loop
|
||
.Lsca_sqrt_end:""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'xmm0', 'rcx', 'r8', 'r10'])
|
||
|
||
|
||
def simd_div_array(a: t.CPtr, b: t.CPtr, out: t.CPtr, n: t.CSizeT):
|
||
"""AVX2 batch: out[i:i+4] = a[i:i+4] / b[i:i+4]"""
|
||
c.Asm(f"""mov rcx, {c.AsmInp(n, t.ASM_DESCR.REG_ANY)}
|
||
shr rcx, 2
|
||
mov r8, {c.AsmInp(a, t.ASM_DESCR.REG_ANY)}
|
||
mov r9, {c.AsmInp(b, t.ASM_DESCR.REG_ANY)}
|
||
mov r10, {c.AsmInp(out, t.ASM_DESCR.REG_ANY)}
|
||
test rcx, rcx
|
||
jz .Ldiv_end
|
||
.Ldiv_loop:
|
||
vmovupd ymm0, [r8]
|
||
vmovupd ymm1, [r9]
|
||
vdivpd ymm0, ymm0, ymm1
|
||
vmovupd [r10], ymm0
|
||
add r8, 32
|
||
add r9, 32
|
||
add r10, 32
|
||
dec rcx
|
||
jnz .Ldiv_loop
|
||
.Ldiv_end:
|
||
vzeroupper""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1', 'rcx', 'r8', 'r9', 'r10'])
|
||
|
||
|
||
def simd_sqrt_array(a: t.CPtr, out: t.CPtr, n: t.CSizeT):
|
||
"""AVX2 batch: out[i:i+4] = sqrt(a[i:i+4])
|
||
×4 展开：4 路独立 vsqrtpd 链隐藏平方根长延迟（吞吐受限场景显著提速）。"""
|
||
c.Asm(f"""mov rcx, {c.AsmInp(n, t.ASM_DESCR.REG_ANY)}
|
||
shr rcx, 2
|
||
mov r8, {c.AsmInp(a, t.ASM_DESCR.REG_ANY)}
|
||
mov r10, {c.AsmInp(out, t.ASM_DESCR.REG_ANY)}
|
||
mov r11, rcx
|
||
shr r11, 2
|
||
test r11, r11
|
||
jz .Ls_sqrt_tail
|
||
.Ls_sqrt_big:
|
||
vmovupd ymm0, [r8]
|
||
vmovupd ymm1, [r8 + 32]
|
||
vmovupd ymm2, [r8 + 64]
|
||
vmovupd ymm3, [r8 + 96]
|
||
vsqrtpd ymm0, ymm0
|
||
vsqrtpd ymm1, ymm1
|
||
vsqrtpd ymm2, ymm2
|
||
vsqrtpd ymm3, ymm3
|
||
vmovupd [r10], ymm0
|
||
vmovupd [r10 + 32], ymm1
|
||
vmovupd [r10 + 64], ymm2
|
||
vmovupd [r10 + 96], ymm3
|
||
add r8, 128
|
||
add r10, 128
|
||
dec r11
|
||
jnz .Ls_sqrt_big
|
||
.Ls_sqrt_tail:
|
||
and rcx, 3
|
||
jz .Ls_sqrt_end
|
||
.Ls_sqrt_tloop:
|
||
vmovupd ymm0, [r8]
|
||
vsqrtpd ymm0, ymm0
|
||
vmovupd [r10], ymm0
|
||
add r8, 32
|
||
add r10, 32
|
||
dec rcx
|
||
jnz .Ls_sqrt_tloop
|
||
.Ls_sqrt_end:
|
||
vzeroupper""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1', 'ymm2', 'ymm3',
|
||
'rcx', 'r8', 'r10', 'r11'])
|
||
|
||
|
||
def simd_abs_array(a: t.CPtr, out: t.CPtr, n: t.CSizeT):
|
||
"""AVX2 batch: out[i:i+4] = |a[i:i+4]|"""
|
||
c.Asm(f"""mov rcx, {c.AsmInp(n, t.ASM_DESCR.REG_ANY)}
|
||
shr rcx, 2
|
||
mov r8, {c.AsmInp(a, t.ASM_DESCR.REG_ANY)}
|
||
mov r10, {c.AsmInp(out, t.ASM_DESCR.REG_ANY)}
|
||
vpcmpeqd ymm2, ymm2, ymm2
|
||
vpsllq ymm2, ymm2, 63
|
||
test rcx, rcx
|
||
jz .Labs_end
|
||
.Labs_loop:
|
||
vmovupd ymm0, [r8]
|
||
vandnpd ymm0, ymm2, ymm0
|
||
vmovupd [r10], ymm0
|
||
add r8, 32
|
||
add r10, 32
|
||
dec rcx
|
||
jnz .Labs_loop
|
||
.Labs_end:
|
||
vzeroupper""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm2', 'rcx', 'r8', 'r10'])
|
||
|
||
|
||
def simd_neg_array(a: t.CPtr, out: t.CPtr, n: t.CSizeT):
|
||
"""AVX2 batch: out[i:i+4] = -a[i:i+4] (flip sign bit via XOR)"""
|
||
c.Asm(f"""mov rcx, {c.AsmInp(n, t.ASM_DESCR.REG_ANY)}
|
||
shr rcx, 2
|
||
mov r8, {c.AsmInp(a, t.ASM_DESCR.REG_ANY)}
|
||
mov r10, {c.AsmInp(out, t.ASM_DESCR.REG_ANY)}
|
||
vpcmpeqd ymm2, ymm2, ymm2
|
||
vpsllq ymm2, ymm2, 63
|
||
test rcx, rcx
|
||
jz .Lneg_end
|
||
.Lneg_loop:
|
||
vmovupd ymm0, [r8]
|
||
vxorpd ymm0, ymm0, ymm2
|
||
vmovupd [r10], ymm0
|
||
add r8, 32
|
||
add r10, 32
|
||
dec rcx
|
||
jnz .Lneg_loop
|
||
.Lneg_end:
|
||
vzeroupper""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm2', 'rcx', 'r8', 'r10'])
|
||
|
||
|
||
def simd_dot_array(a: t.CPtr, b: t.CPtr, n: t.CSizeT, out: t.CPtr):
|
||
"""AVX2 batch: out[0] = dot(a[0:n], b[0:n]) using FMA accumulation"""
|
||
c.Asm(f"""mov rcx, {c.AsmInp(n, t.ASM_DESCR.REG_ANY)}
|
||
shr rcx, 2
|
||
mov r8, {c.AsmInp(a, t.ASM_DESCR.REG_ANY)}
|
||
mov r9, {c.AsmInp(b, t.ASM_DESCR.REG_ANY)}
|
||
vxorpd ymm2, ymm2, ymm2
|
||
test rcx, rcx
|
||
jz .Ldot_end
|
||
.Ldot_loop:
|
||
vmovupd ymm0, [r8]
|
||
vmovupd ymm1, [r9]
|
||
vfmadd231pd ymm2, ymm0, ymm1
|
||
add r8, 32
|
||
add r9, 32
|
||
dec rcx
|
||
jnz .Ldot_loop
|
||
.Ldot_end:
|
||
vextractf128 xmm1, ymm2, 1
|
||
vaddpd xmm2, xmm2, xmm1
|
||
vhaddpd xmm2, xmm2, xmm2
|
||
vmovsd [{c.AsmInp(out, t.ASM_DESCR.REG_ANY)}], xmm2
|
||
vzeroupper""",
|
||
op=[t.ASM_DESCR.CLOBBER_MEMORY, 'ymm0', 'ymm1', 'ymm2', 'rcx', 'r8', 'r9'])
|
||
|
||
|
||
# ============================================================
|
||
# LLVMIR precise scalar operations (LLVM intrinsic functions)
|
||
# ============================================================
|
||
|
||
def simd_sqrt(x: t.CDouble) -> t.CDouble:
|
||
"""LLVM intrinsic: sqrt(x)"""
|
||
out: t.CDouble
|
||
c.LLVMIR(f"{c.LOut(out)} = call double @llvm.sqrt.f64(double {c.LInp(x)})")
|
||
return out
|
||
|
||
|
||
def simd_fabs(x: t.CDouble) -> t.CDouble:
|
||
"""LLVM intrinsic: |x|"""
|
||
out: t.CDouble
|
||
c.LLVMIR(f"{c.LOut(out)} = call double @llvm.fabs.f64(double {c.LInp(x)})")
|
||
return out
|
||
|
||
|
||
def simd_floor(x: t.CDouble) -> t.CDouble:
|
||
"""LLVM intrinsic: floor(x)"""
|
||
out: t.CDouble
|
||
c.LLVMIR(f"{c.LOut(out)} = call double @llvm.floor.f64(double {c.LInp(x)})")
|
||
return out
|
||
|
||
|
||
def simd_ceil(x: t.CDouble) -> t.CDouble:
|
||
"""LLVM intrinsic: ceil(x)"""
|
||
out: t.CDouble
|
||
c.LLVMIR(f"{c.LOut(out)} = call double @llvm.ceil.f64(double {c.LInp(x)})")
|
||
return out
|
||
|
||
|
||
def simd_round(x: t.CDouble) -> t.CDouble:
|
||
"""LLVM intrinsic: round(x)"""
|
||
out: t.CDouble
|
||
c.LLVMIR(f"{c.LOut(out)} = call double @llvm.round.f64(double {c.LInp(x)})")
|
||
return out
|
||
|
||
|
||
def simd_trunc(x: t.CDouble) -> t.CDouble:
|
||
"""LLVM intrinsic: trunc(x)"""
|
||
out: t.CDouble
|
||
c.LLVMIR(f"{c.LOut(out)} = call double @llvm.trunc.f64(double {c.LInp(x)})")
|
||
return out
|
||
|
||
|
||
def simd_fma(a: t.CDouble, b: t.CDouble, c_val: t.CDouble) -> t.CDouble:
|
||
"""LLVM intrinsic: a*b + c (single rounding)"""
|
||
out: t.CDouble
|
||
c.LLVMIR(f"{c.LOut(out)} = call double @llvm.fma.f64(double {c.LInp(a)}, double {c.LInp(b)}, double {c.LInp(c_val)})")
|
||
return out
|
||
|
||
|
||
def simd_copysign(mag: t.CDouble, sign: t.CDouble) -> t.CDouble:
|
||
"""LLVM intrinsic: copysign(mag, sign)"""
|
||
out: t.CDouble
|
||
c.LLVMIR(f"{c.LOut(out)} = call double @llvm.copysign.f64(double {c.LInp(mag)}, double {c.LInp(sign)})")
|
||
return out
|
||
|
||
|
||
def simd_neg(x: t.CDouble) -> t.CDouble:
|
||
"""LLVM intrinsic: -x (via copysign, correct for negative/-0.0/NaN inputs)
|
||
|
||
copysign(x, -1.0) is wrong for x < 0 (returns -|x| instead of +|x|).
|
||
copysign(x, -x) = |x| * sign(-x) = -x is correct for all IEEE 754 values.
|
||
"""
|
||
out: t.CDouble
|
||
c.LLVMIR(f"{c.LOut(out)} = call double @llvm.copysign.f64(double {c.LInp(x)}, double {c.LInp(t.CDouble(0.0) - x)})")
|
||
return out
|
||
|
||
|
||
def simd_minnum(a: t.CDouble, b: t.CDouble) -> t.CDouble:
|
||
"""LLVM intrinsic: min(a, b) (NaN-aware)"""
|
||
out: t.CDouble
|
||
c.LLVMIR(f"{c.LOut(out)} = call double @llvm.minnum.f64(double {c.LInp(a)}, double {c.LInp(b)})")
|
||
return out
|
||
|
||
|
||
def simd_maxnum(a: t.CDouble, b: t.CDouble) -> t.CDouble:
|
||
"""LLVM intrinsic: max(a, b) (NaN-aware)"""
|
||
out: t.CDouble
|
||
c.LLVMIR(f"{c.LOut(out)} = call double @llvm.maxnum.f64(double {c.LInp(a)}, double {c.LInp(b)})")
|
||
return out
|
||
|
||
|
||
def simd_exp(x: t.CDouble) -> t.CDouble:
|
||
"""LLVM intrinsic: e^x"""
|
||
out: t.CDouble
|
||
c.LLVMIR(f"{c.LOut(out)} = call double @llvm.exp.f64(double {c.LInp(x)})")
|
||
return out
|
||
|
||
|
||
def simd_log(x: t.CDouble) -> t.CDouble:
|
||
"""LLVM intrinsic: ln(x)"""
|
||
out: t.CDouble
|
||
c.LLVMIR(f"{c.LOut(out)} = call double @llvm.log.f64(double {c.LInp(x)})")
|
||
return out
|
||
|
||
|
||
def simd_sin(x: t.CDouble) -> t.CDouble:
|
||
"""LLVM intrinsic: sin(x)"""
|
||
out: t.CDouble
|
||
c.LLVMIR(f"{c.LOut(out)} = call double @llvm.sin.f64(double {c.LInp(x)})")
|
||
return out
|
||
|
||
|
||
def simd_cos(x: t.CDouble) -> t.CDouble:
|
||
"""LLVM intrinsic: cos(x)"""
|
||
out: t.CDouble
|
||
c.LLVMIR(f"{c.LOut(out)} = call double @llvm.cos.f64(double {c.LInp(x)})")
|
||
return out
|
||
|
||
|
||
def simd_pow(x: t.CDouble, y: t.CDouble) -> t.CDouble:
|
||
"""LLVM intrinsic: x^y"""
|
||
out: t.CDouble
|
||
c.LLVMIR(f"{c.LOut(out)} = call double @llvm.pow.f64(double {c.LInp(x)}, double {c.LInp(y)})")
|
||
return out
|