gather
diff --git a/src/kernel.rs b/src/kernel.rs index 4ac76eb..cf67bf5 100644 --- a/src/kernel.rs +++ b/src/kernel.rs
@@ -827,113 +827,59 @@ // This routine loads and transposes message words, populates the rest of the state registers, // and invokes blake3_avx512_kernel_16. // -------------------------------------------------------------------------------------------- + ".p2align 6", + "BLAKE3_CHUNK_OFFSETS_16:", + ".long 0, 1024, 2048, 3072, 4096, 5120, 6144, 7168, 8192, 9216, 10240, 11264, 12288, 13312, 14336, 15360", "blake3_avx512_blocks_16:", - // Load the message blocks first (unaligned). See the comments immediately below for why we - // choose these registers. - "vmovdqu32 zmm24, zmmword ptr [rdi + 0 * 1024]", - "vmovdqu32 zmm25, zmmword ptr [rdi + 1 * 1024]", - "vmovdqu32 zmm26, zmmword ptr [rdi + 2 * 1024]", - "vmovdqu32 zmm27, zmmword ptr [rdi + 3 * 1024]", - "vmovdqu32 zmm28, zmmword ptr [rdi + 4 * 1024]", - "vmovdqu32 zmm29, zmmword ptr [rdi + 5 * 1024]", - "vmovdqu32 zmm30, zmmword ptr [rdi + 6 * 1024]", - "vmovdqu32 zmm31, zmmword ptr [rdi + 7 * 1024]", - "vmovdqu32 zmm8, zmmword ptr [rdi + 8 * 1024]", - "vmovdqu32 zmm9, zmmword ptr [rdi + 9 * 1024]", - "vmovdqu32 zmm10, zmmword ptr [rdi + 10 * 1024]", - "vmovdqu32 zmm11, zmmword ptr [rdi + 11 * 1024]", - "vmovdqu32 zmm12, zmmword ptr [rdi + 12 * 1024]", - "vmovdqu32 zmm13, zmmword ptr [rdi + 13 * 1024]", - "vmovdqu32 zmm14, zmmword ptr [rdi + 14 * 1024]", - "vmovdqu32 zmm15, zmmword ptr [rdi + 15 * 1024]", - // Transpose the message blocks. This requires a few different passes: - // 1) interleave 32-bit lanes - // 2) interleave 64-bit lanes - // 3) interleave 128-bit lanes - // 4) interleave 256-bit lanes (but there's no such instruction, so actually 128 bits again) - // The last of these passes is easier to implement if we can make use of 8 scratch registers. - // zmm0-zmm7 are holding the incoming CV and we don't want to touch those. But zmm8-zmm15 - // aren't holding anything important, and we can use those as long as we reinitialize them - // before we run the kernel. For consistency, we'll use all 8 scratch registers for each pass - // (even though the earlier passes would be fine using fewer), and we'll have each pass rotate - // our 24 message+scratch vectors 8 places "to the left". Thus starting 8 places "to the right" - // in the rotation lets us end up on target after 4 passes, and that's why we loaded the - // message vectors in the order we did above. - // - // The first pass, interleaving 32-bit lanes. Here's the first vector before: - // (zmm24) a0, a1, a2, a3, a4, a5, a8, a7, a8, a9, a10, a11, a12, a13, a14, a15 - // And after: - // (zmm16) a0, b0, a1, b1, a4, b4, a5, b5, a8, b8, a9, b9, a12, b12, a13, b13 - "vpunpckldq zmm16, zmm24, zmm25", - "vpunpckhdq zmm17, zmm24, zmm25", - "vpunpckldq zmm18, zmm26, zmm27", - "vpunpckhdq zmm19, zmm26, zmm27", - "vpunpckldq zmm20, zmm28, zmm29", - "vpunpckhdq zmm21, zmm28, zmm29", - "vpunpckldq zmm22, zmm30, zmm31", - "vpunpckhdq zmm23, zmm30, zmm31", - "vpunpckldq zmm24, zmm8, zmm9", - "vpunpckhdq zmm25, zmm8, zmm9", - "vpunpckldq zmm26, zmm10, zmm11", - "vpunpckhdq zmm27, zmm10, zmm11", - "vpunpckldq zmm28, zmm12, zmm13", - "vpunpckhdq zmm29, zmm12, zmm13", - "vpunpckldq zmm30, zmm14, zmm15", - "vpunpckhdq zmm31, zmm14, zmm15", - // The second pass, interleaving 64-bit lanes. After this the first vector will be: - // (zmm8) a0, b0, c0, d0, a4, b4, c4, d4, a8, b8, c8, d8, a12, b12, c12, d12 - "vpunpcklqdq zmm8, zmm16, zmm18", - "vpunpckhqdq zmm9, zmm16, zmm18", - "vpunpcklqdq zmm10, zmm17, zmm19", - "vpunpckhqdq zmm11, zmm17, zmm19", - "vpunpcklqdq zmm12, zmm20, zmm22", - "vpunpckhqdq zmm13, zmm20, zmm22", - "vpunpcklqdq zmm14, zmm21, zmm23", - "vpunpckhqdq zmm15, zmm21, zmm23", - "vpunpcklqdq zmm16, zmm24, zmm26", - "vpunpckhqdq zmm17, zmm24, zmm26", - "vpunpcklqdq zmm18, zmm25, zmm27", - "vpunpckhqdq zmm19, zmm25, zmm27", - "vpunpcklqdq zmm20, zmm28, zmm30", - "vpunpckhqdq zmm21, zmm28, zmm30", - "vpunpcklqdq zmm22, zmm29, zmm31", - "vpunpckhqdq zmm23, zmm29, zmm31", - // The third pass, interleaving 128-bit lanes. After this the first vector will be: - // (zmm24) a0, b0, c0, d0, a8, b8, c8, d8, e0, f0, g0, h0, e8, f8, g8, h8 - "vshufi32x4 zmm24, zmm8, zmm12, 0x88", // 0b10001000: lo 128-bit lanes A0/A2/B0/B2 - "vshufi32x4 zmm25, zmm9, zmm13, 0x88", - "vshufi32x4 zmm26, zmm10, zmm14, 0x88", - "vshufi32x4 zmm27, zmm11, zmm15, 0x88", - "vshufi32x4 zmm28, zmm8, zmm12, 0xdd", // 0b11011101: hi 128-bit lanes A1/A3/B1/B3 - "vshufi32x4 zmm29, zmm9, zmm13, 0xdd", - "vshufi32x4 zmm30, zmm10, zmm14, 0xdd", - "vshufi32x4 zmm31, zmm11, zmm15, 0xdd", - "vshufi32x4 zmm8, zmm16, zmm20, 0x88", // lo - "vshufi32x4 zmm9, zmm17, zmm21, 0x88", - "vshufi32x4 zmm10, zmm18, zmm22, 0x88", - "vshufi32x4 zmm11, zmm19, zmm23, 0x88", - "vshufi32x4 zmm12, zmm16, zmm20, 0xdd", // hi - "vshufi32x4 zmm13, zmm17, zmm21, 0xdd", - "vshufi32x4 zmm14, zmm18, zmm22, 0xdd", - "vshufi32x4 zmm15, zmm19, zmm23, 0xdd", - // The fourth and final pass, interleaving 128-bit lanes again. The first vector will be: - // (zmm16) a0, b0, c0, d0, e0, f0, g0, h0, i0, j0, k0, l0, m0, n0, o0, p0 - "vshufi32x4 zmm16, zmm24, zmm8, 0x88", // lo - "vshufi32x4 zmm17, zmm25, zmm9, 0x88", - "vshufi32x4 zmm18, zmm26, zmm10, 0x88", - "vshufi32x4 zmm19, zmm27, zmm11, 0x88", - "vshufi32x4 zmm20, zmm28, zmm12, 0x88", - "vshufi32x4 zmm21, zmm29, zmm13, 0x88", - "vshufi32x4 zmm22, zmm30, zmm14, 0x88", - "vshufi32x4 zmm23, zmm31, zmm15, 0x88", - "vshufi32x4 zmm24, zmm24, zmm8, 0xdd", // hi - "vshufi32x4 zmm25, zmm25, zmm9, 0xdd", - "vshufi32x4 zmm26, zmm26, zmm10, 0xdd", - "vshufi32x4 zmm27, zmm27, zmm11, 0xdd", - "vshufi32x4 zmm28, zmm28, zmm12, 0xdd", - "vshufi32x4 zmm29, zmm29, zmm13, 0xdd", - "vshufi32x4 zmm30, zmm30, zmm14, 0xdd", - "vshufi32x4 zmm31, zmm31, zmm15, 0xdd", + "vmovdqa32 zmm8, zmmword ptr [rip + BLAKE3_CHUNK_OFFSETS_16]", + "kxnorw k1, k0, k0", + "vpxorq xmm16, xmm16, xmm16", + "vpgatherdd zmm16{{k1}}, [rdi + ( 0 * 4) + zmm8]", + "kxnorw k1, k0, k0", + "vpxorq xmm17, xmm17, xmm17", + "vpgatherdd zmm17{{k1}}, [rdi + ( 1 * 4) + zmm8]", + "kxnorw k1, k0, k0", + "vpxorq xmm18, xmm18, xmm18", + "vpgatherdd zmm18{{k1}}, [rdi + ( 2 * 4) + zmm8]", + "kxnorw k1, k0, k0", + "vpxorq xmm19, xmm19, xmm19", + "vpgatherdd zmm19{{k1}}, [rdi + ( 3 * 4) + zmm8]", + "kxnorw k1, k0, k0", + "vpxorq xmm20, xmm20, xmm20", + "vpgatherdd zmm20{{k1}}, [rdi + ( 4 * 4) + zmm8]", + "kxnorw k1, k0, k0", + "vpxorq xmm21, xmm21, xmm21", + "vpgatherdd zmm21{{k1}}, [rdi + ( 5 * 4) + zmm8]", + "kxnorw k1, k0, k0", + "vpxorq xmm22, xmm22, xmm22", + "vpgatherdd zmm22{{k1}}, [rdi + ( 6 * 4) + zmm8]", + "kxnorw k1, k0, k0", + "vpxorq xmm23, xmm23, xmm23", + "vpgatherdd zmm23{{k1}}, [rdi + ( 7 * 4) + zmm8]", + "kxnorw k1, k0, k0", + "vpxorq xmm24, xmm24, xmm24", + "vpgatherdd zmm24{{k1}}, [rdi + ( 8 * 4) + zmm8]", + "kxnorw k1, k0, k0", + "vpxorq xmm25, xmm25, xmm25", + "vpgatherdd zmm25{{k1}}, [rdi + ( 9 * 4) + zmm8]", + "kxnorw k1, k0, k0", + "vpxorq xmm26, xmm26, xmm26", + "vpgatherdd zmm26{{k1}}, [rdi + (10 * 4) + zmm8]", + "kxnorw k1, k0, k0", + "vpxorq xmm27, xmm27, xmm27", + "vpgatherdd zmm27{{k1}}, [rdi + (11 * 4) + zmm8]", + "kxnorw k1, k0, k0", + "vpxorq xmm28, xmm28, xmm28", + "vpgatherdd zmm28{{k1}}, [rdi + (12 * 4) + zmm8]", + "kxnorw k1, k0, k0", + "vpxorq xmm29, xmm29, xmm29", + "vpgatherdd zmm29{{k1}}, [rdi + (13 * 4) + zmm8]", + "kxnorw k1, k0, k0", + "vpxorq xmm30, xmm30, xmm30", + "vpgatherdd zmm30{{k1}}, [rdi + (14 * 4) + zmm8]", + "kxnorw k1, k0, k0", + "vpxorq xmm31, xmm31, xmm31", + "vpgatherdd zmm31{{k1}}, [rdi + (15 * 4) + zmm8]", // Initialize the third and fourth rows of the state, which we just used as scratch space // during transposition. "vmovdqa32 zmm8, zmmword ptr [BLAKE3_IV0_16 + rip]", // IV constants