fix: Fix asm_decompose

sp301415 · Jan 26, 2025 · 590b9b2 · 590b9b2
1 parent 8492f3c
commit 590b9b2
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 14 deletions.
diff --git a/internal/asmgen/decompose.go b/internal/asmgen/decompose.go
@@ -79,9 +79,21 @@ func decomposePolyAssignUint32AVX2() {
 	SUBQ(Imm(3), jj)
 
 	Label("level_loop_end")
-	CMPQ(j, Imm(0))
+	CMPQ(j, Imm(1))
 	JGE(LabelRef("level_loop_body"))
 
+	u = YMM()
+	VANDPD(baseMask, c, u)
+
+	uCarry = YMM()
+	VANDPD(baseHalf, u, uCarry)
+	VPSLLD(Imm(1), uCarry, uCarry)
+	VPSUBD(uCarry, u, u)
+
+	decomposedOut0 := GP64()
+	MOVQ(Mem{Base: decomposedOut}, decomposedOut0)
+	VMOVDQU(u, Mem{Base: decomposedOut0, Index: i, Scale: 4})
+
 	ADDQ(Imm(8), i)
 
 	Label("N_loop_end")
@@ -159,9 +171,21 @@ func decomposePolyAssignUint64AVX2() {
 	SUBQ(Imm(3), jj)
 
 	Label("level_loop_end")
-	CMPQ(j, Imm(0))
+	CMPQ(j, Imm(1))
 	JGE(LabelRef("level_loop_body"))
 
+	u = YMM()
+	VANDPD(baseMask, c, u)
+
+	uCarry = YMM()
+	VANDPD(baseHalf, u, uCarry)
+	VPSLLQ(Imm(1), uCarry, uCarry)
+	VPSUBQ(uCarry, u, u)
+
+	decomposedOut0 := GP64()
+	MOVQ(Mem{Base: decomposedOut}, decomposedOut0)
+	VMOVDQU(u, Mem{Base: decomposedOut0, Index: i, Scale: 8})
+
 	ADDQ(Imm(4), i)
 
 	Label("N_loop_end")

diff --git a/math/poly/asm_convert.go b/math/poly/asm_convert.go
@@ -92,9 +92,7 @@ func convertPolyToFourierPolyAssign[T num.Integer](p []T, fpOut []float64) {
 // floatModQInPlace computes coeffs mod Q in place.
 func floatModQInPlace(coeffs []float64, Q float64) {
 	for i := range coeffs {
-		cQuo := coeffs[i] / Q
-		cRem := cQuo - math.Round(cQuo)
-		coeffs[i] = math.Round(cRem * Q)
+		coeffs[i] = math.Round(coeffs[i] - Q*math.Round(coeffs[i]/Q))
 	}
 }
 

diff --git a/math/poly/asm_convert_amd64.go b/math/poly/asm_convert_amd64.go
@@ -116,9 +116,7 @@ func floatModQInPlace(coeffs []float64, Q float64) {
 	}
 
 	for i := range coeffs {
-		cQuo := coeffs[i] / Q
-		cRem := cQuo - math.Round(cQuo)
-		coeffs[i] = math.Round(cRem * Q)
+		coeffs[i] = math.Round(coeffs[i] - Q*math.Round(coeffs[i]/Q))
 	}
 }
 

diff --git a/tfhe/asm_decompose_amd64.s b/tfhe/asm_decompose_amd64.s
@@ -52,9 +52,15 @@ level_loop_body:
 	SUBQ    $0x03, R8
 
 level_loop_end:
-	CMPQ DI, $0x00
-	JGE  level_loop_body
-	ADDQ $0x08, SI
+	CMPQ    DI, $0x01
+	JGE     level_loop_body
+	VANDPD  Y4, Y6, Y6
+	VANDPD  Y0, Y6, Y7
+	VPSLLD  $0x01, Y7, Y7
+	VPSUBD  Y7, Y6, Y6
+	MOVQ    (CX), DI
+	VMOVDQU Y6, (DI)(SI*4)
+	ADDQ    $0x08, SI
 
 N_loop_end:
 	CMPQ SI, DX
@@ -106,9 +112,15 @@ level_loop_body:
 	SUBQ    $0x03, R8
 
 level_loop_end:
-	CMPQ DI, $0x00
-	JGE  level_loop_body
-	ADDQ $0x04, SI
+	CMPQ    DI, $0x01
+	JGE     level_loop_body
+	VANDPD  Y4, Y6, Y6
+	VANDPD  Y0, Y6, Y7
+	VPSLLQ  $0x01, Y7, Y7
+	VPSUBQ  Y7, Y6, Y6
+	MOVQ    (CX), DI
+	VMOVDQU Y6, (DI)(SI*8)
+	ADDQ    $0x04, SI
 
 N_loop_end:
 	CMPQ SI, DX