makarna/pkg/backend/cpu/nn/silu_avx512.s

//go:build amd64
// +build amd64

#include "textflag.h"

// func siluAVX512Asm(x *float32, n int)
TEXT ·siluAVX512Asm(SB), NOSPLIT, $0-16
	// Load args
	MOVQ x+0(FP), DI
	MOVQ n+8(FP), CX

	CMPQ CX, $0
	JLE done

	// Broadcast constants
	VBROADCASTSS ·expHi(SB), Z14
	VBROADCASTSS ·expLo(SB), Z13
	VBROADCASTSS ·log2EF(SB), Z12
	VBROADCASTSS ·halfConst(SB), Z11
	VBROADCASTSS ·expC1(SB), Z10
	VBROADCASTSS ·expC2(SB), Z9
	VBROADCASTSS ·oneConst(SB), Z8
	VPBROADCASTD ·signMaskConst(SB), Z15

loop:
	CMPQ CX, $16
	JL done

	VMOVUPS (DI), Z0         // original x
	VMOVAPS Z0, Z1           // copy for neg
	VXORPS Z15, Z1, Z1       // z1 = -x

	VMINPS Z14, Z1, Z1       // clamp hi
	VMAXPS Z13, Z1, Z1       // clamp lo

	VMULPS Z12, Z1, Z2       // z2 = x * log2e
	VADDPS Z11, Z2, Z2       // +0.5
	VRNDSCALEPS $1, Z2, Z2   // floor

	VCVTPS2DQ Z2, Z6         // integer exponent
	VCVTDQ2PS Z6, Z5         // fx as float

	VMULPS Z10, Z5, Z3       // fx * C1
	VSUBPS Z3, Z1, Z1
	VMULPS Z9, Z5, Z3        // fx * C2
	VSUBPS Z3, Z1, Z1

	VMULPS Z1, Z1, Z3        // z = x*x

	VBROADCASTSS ·polyP0(SB), Z4
	VMULPS Z1, Z4, Z4
	VBROADCASTSS ·polyP1(SB), Z5
	VADDPS Z5, Z4, Z4
	VMULPS Z1, Z4, Z4
	VBROADCASTSS ·polyP2(SB), Z5
	VADDPS Z5, Z4, Z4
	VMULPS Z1, Z4, Z4
	VBROADCASTSS ·polyP3(SB), Z5
	VADDPS Z5, Z4, Z4
	VMULPS Z1, Z4, Z4
	VBROADCASTSS ·polyP4(SB), Z5
	VADDPS Z5, Z4, Z4
	VMULPS Z1, Z4, Z4
	VBROADCASTSS ·polyP5(SB), Z5
	VADDPS Z5, Z4, Z4

	VMULPS Z3, Z4, Z4        // y *= z
	VADDPS Z1, Z4, Z4        // y += x
	VADDPS Z8, Z4, Z4        // y += 1

	VPBROADCASTD ·expBiasConst(SB), Z5
	VPADDD Z5, Z6, Z6
	VPSLLD $23, Z6, Z6
	VMULPS Z6, Z4, Z4        // exp(-x)

	VADDPS Z8, Z4, Z3        // denom = 1 + exp(-x)
	VDIVPS Z3, Z8, Z3        // 1 / denom
	VMULPS Z0, Z3, Z0        // x * sigmoid(x)
	VMOVUPS Z0, (DI)

	ADDQ $64, DI
	SUBQ $16, CX
	JMP loop

done:
	RET