| #! /usr/bin/env python3 |
| |
| # Generate asm! |
| # |
| # TODOs: |
| # - vzeroupper |
| # - CET |
| # - prefetches |
| |
| from dataclasses import dataclass, replace |
| |
| X86_64 = "x86_64" |
| AVX512 = "avx512" |
| AVX2 = "avx2" |
| SSE41 = "sse41" |
| SSE2 = "sse2" |
| LINUX = "linux" |
| |
| MESSAGE_SCHEDULE = [ |
| [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15], |
| [2, 6, 3, 10, 7, 0, 4, 13, 1, 11, 12, 5, 9, 14, 15, 8], |
| [3, 4, 10, 12, 13, 2, 7, 14, 6, 5, 9, 0, 11, 15, 8, 1], |
| [10, 7, 12, 9, 14, 3, 13, 15, 4, 0, 11, 2, 5, 8, 1, 6], |
| [12, 13, 9, 11, 15, 10, 14, 8, 7, 2, 5, 3, 0, 1, 6, 4], |
| [9, 14, 11, 5, 8, 12, 15, 1, 13, 3, 0, 10, 2, 6, 4, 7], |
| [11, 15, 5, 0, 1, 9, 8, 6, 14, 10, 2, 12, 3, 4, 7, 13], |
| ] |
| |
| |
| @dataclass |
| class Target: |
| arch: str |
| extension: str |
| os: str |
| |
| def arg32(self, index): |
| system_v_args_32 = ["edi", "esi", "edx", "ecx", "r8d", "r9d"] |
| return system_v_args_32[index] |
| |
| def arg64(self, index): |
| system_v_args_64 = ["rdi", "rsi", "rdx", "rcx", "r8", "r9"] |
| return system_v_args_64[index] |
| |
| def scratch32(self, index): |
| system_v_scratch_32 = ["eax", "r10d", "r11d"] |
| return system_v_scratch_32[index] |
| |
| def scratch64(self, index): |
| system_v_scratch_64 = ["rax", "r10", "r11"] |
| return system_v_scratch_64[index] |
| |
| def reg128(self, index): |
| assert self.arch == X86_64 |
| return "xmm" + str(index) |
| |
| def ret(self): |
| return "ret" |
| |
| |
| def add_row(target, output, degree, dest, src): |
| assert target.arch == X86_64 |
| if target.extension == AVX512: |
| if degree == 1: |
| output.append(f"vpaddd xmm{dest}, xmm{dest}, xmm{src}") |
| elif degree in (2, 8): |
| output.append(f"vpaddd ymm{dest}, ymm{dest}, ymm{src}") |
| elif degree in (4, 16): |
| output.append(f"vpaddd zmm{dest}, zmm{dest}, zmm{src}") |
| else: |
| raise NotImplementedError |
| elif target.extension == AVX2: |
| assert degree == 2 |
| output.append(f"vpaddd ymm{dest}, ymm{dest}, ymm{src}") |
| elif target.extension in (SSE41, SSE2): |
| assert degree == 1 |
| output.append(f"paddd xmm{dest}, xmm{src}") |
| else: |
| raise NotImplementedError |
| |
| |
| def xor_row(target, output, degree, dest, src): |
| assert target.arch == X86_64 |
| if target.extension == AVX512: |
| if degree == 1: |
| output.append(f"vpxord xmm{dest}, xmm{dest}, xmm{src}") |
| elif degree in (2, 8): |
| output.append(f"vpxord ymm{dest}, ymm{dest}, ymm{src}") |
| elif degree in (4, 16): |
| output.append(f"vpxord zmm{dest}, zmm{dest}, zmm{src}") |
| else: |
| raise NotImplementedError |
| elif target.extension == AVX2: |
| assert degree == 2 |
| output.append(f"vpxord ymm{dest}, ymm{dest}, ymm{src}") |
| elif target.extension in (SSE41, SSE2): |
| assert degree == 1 |
| output.append(f"pxor xmm{dest}, xmm{src}") |
| else: |
| raise NotImplementedError |
| |
| |
| # This is the >>> operation in G, not state diagonalization or message permutation. |
| def bitrotate_row(target, output, degree, reg, bits): |
| assert target.arch == X86_64 |
| if target.extension == AVX512: |
| if degree == 1: |
| output.append(f"vprord xmm{reg}, xmm{reg}, {bits}") |
| elif degree in (2, 8): |
| output.append(f"vprord ymm{reg}, ymm{reg}, {bits}") |
| elif degree in (4, 16): |
| output.append(f"vprord zmm{reg}, zmm{reg}, {bits}") |
| else: |
| raise NotImplementedError |
| elif target.extension == AVX2: |
| assert degree == 2 |
| if bits == 16: |
| # ymm14 is initialized at the start of the kernel |
| output.append(f"vpshufb ymm{reg}, ymm{reg}, ymm14") |
| elif bits == 8: |
| # ymm15 is initialized at the start of the kernel |
| output.append(f"vpshufb ymm{reg}, ymm{reg}, ymm15") |
| else: |
| # Do two bitshifts, using register 8 as temp. |
| output.append(f"vpsrld ymm8, ymm{reg}, {bits}") |
| output.append(f"vpslld ymm{reg}, ymm{reg}, {32 - bits}") |
| output.append(f"vpor ymm{reg}, ymm{reg}, ymm8") |
| elif target.extension == SSE41: |
| assert degree == 1 |
| if bits == 16: |
| # xmm15 is initialized at the start of the kernel |
| output.append(f"pshufb xmm{reg}, xmm15") |
| elif bits == 8: |
| # xmm14 is initialized at the start of the kernel |
| output.append(f"pshufb xmm{reg}, xmm14") |
| else: |
| # Do two bitshifts, using register 11 as temp. |
| output.append(f"movdqa xmm11, xmm{reg}") |
| output.append(f"pslld xmm{reg}, {32 - bits}") |
| output.append(f"psrld xmm11, {bits}") |
| output.append(f"por xmm{reg}, xmm11") |
| elif target.extension == SSE2: |
| assert degree == 1 |
| if bits == 16: |
| output.append(f"pshuflw xmm{reg}, xmm{reg}, 0xB1") |
| output.append(f"pshufhw xmm{reg}, xmm{reg}, 0xB1") |
| else: |
| # Do two bitshifts, using register 11 as temp. |
| output.append(f"movdqa xmm11, xmm{reg}") |
| output.append(f"pslld xmm{reg}, {32 - bits}") |
| output.append(f"psrld xmm11, {bits}") |
| output.append(f"por xmm{reg}, xmm11") |
| else: |
| raise NotImplementedError |
| |
| |
| # See the comments above kernel_2d(). |
| def diagonalize_state_rows(target, output, degree): |
| if target.extension == AVX512: |
| if degree == 1: |
| output.append("vpshufd xmm0, xmm0, 0x93") # 3 0 1 2 |
| output.append("vpshufd xmm3, xmm3, 0x4E") # 2 3 0 1 |
| output.append("vpshufd xmm2, xmm2, 0x39") # 1 2 3 0 |
| elif degree == 2: |
| output.append("vpshufd ymm0, ymm0, 0x93") |
| output.append("vpshufd ymm3, ymm3, 0x4E") |
| output.append("vpshufd ymm2, ymm2, 0x39") |
| elif degree == 4: |
| output.append("vpshufd zmm0, zmm0, 0x93") |
| output.append("vpshufd zmm3, zmm3, 0x4E") |
| output.append("vpshufd zmm2, zmm2, 0x39") |
| else: |
| raise NotImplementedError |
| elif target.extension == AVX2: |
| assert degree == 2 |
| output.append("vpshufd ymm0, ymm0, 0x93") |
| output.append("vpshufd ymm3, ymm3, 0x4E") |
| output.append("vpshufd ymm2, ymm2, 0x39") |
| elif target.extension in (SSE41, SSE2): |
| assert degree == 1 |
| output.append("pshufd xmm0, xmm0, 0x93") |
| output.append("pshufd xmm3, xmm3, 0x4E") |
| output.append("pshufd xmm2, xmm2, 0x39") |
| else: |
| raise NotImplementedError |
| |
| |
| # See the comments above kernel_2d(). |
| def undiagonalize_state_rows(target, output, degree): |
| if target.extension == AVX512: |
| if degree == 1: |
| output.append("vpshufd xmm0, xmm0, 0x39") # 1 2 3 0 |
| output.append("vpshufd xmm3, xmm3, 0x4E") # 2 3 0 1 |
| output.append("vpshufd xmm2, xmm2, 0x93") # 3 0 1 2 |
| elif degree == 2: |
| output.append("vpshufd ymm0, ymm0, 0x39") |
| output.append("vpshufd ymm3, ymm3, 0x4E") |
| output.append("vpshufd ymm2, ymm2, 0x93") |
| elif degree == 4: |
| output.append("vpshufd zmm0, zmm0, 0x39") |
| output.append("vpshufd zmm3, zmm3, 0x4E") |
| output.append("vpshufd zmm2, zmm2, 0x93") |
| else: |
| raise NotImplementedError |
| elif target.extension == AVX2: |
| assert degree == 2 |
| output.append("vpshufd ymm0, ymm0, 0x39") |
| output.append("vpshufd ymm3, ymm3, 0x4E") |
| output.append("vpshufd ymm2, ymm2, 0x93") |
| elif target.extension in (SSE41, SSE2): |
| assert degree == 1 |
| output.append("pshufd xmm0, xmm0, 0x39") |
| output.append("pshufd xmm3, xmm3, 0x4E") |
| output.append("pshufd xmm2, xmm2, 0x93") |
| else: |
| raise NotImplementedError |
| |
| |
| # See the comments above kernel_2d(). |
| def permute_message_rows(target, output, degree): |
| if target.extension == AVX512: |
| if degree == 1: |
| output.append("vshufps xmm8, xmm4, xmm5, 214") |
| output.append("vpshufd xmm9, xmm4, 0x0F") |
| output.append("vpshufd xmm4, xmm8, 0x39") |
| output.append("vshufps xmm8, xmm6, xmm7, 250") |
| output.append("vpblendd xmm9, xmm9, xmm8, 0xAA") |
| output.append("vpunpcklqdq xmm8, xmm7, xmm5") |
| output.append("vpblendd xmm8, xmm8, xmm6, 0x88") |
| output.append("vpshufd xmm8, xmm8, 0x78") |
| output.append("vpunpckhdq xmm5, xmm5, xmm7") |
| output.append("vpunpckldq xmm6, xmm6, xmm5") |
| output.append("vpshufd xmm7, xmm6, 0x1E") |
| output.append("vmovdqa xmm5, xmm9") |
| output.append("vmovdqa xmm6, xmm8") |
| elif degree == 2: |
| output.append("vshufps ymm8, ymm4, ymm5, 214") |
| output.append("vpshufd ymm9, ymm4, 0x0F") |
| output.append("vpshufd ymm4, ymm8, 0x39") |
| output.append("vshufps ymm8, ymm6, ymm7, 250") |
| output.append("vpblendd ymm9, ymm9, ymm8, 0xAA") |
| output.append("vpunpcklqdq ymm8, ymm7, ymm5") |
| output.append("vpblendd ymm8, ymm8, ymm6, 0x88") |
| output.append("vpshufd ymm8, ymm8, 0x78") |
| output.append("vpunpckhdq ymm5, ymm5, ymm7") |
| output.append("vpunpckldq ymm6, ymm6, ymm5") |
| output.append("vpshufd ymm7, ymm6, 0x1E") |
| output.append("vmovdqa ymm5, ymm9") |
| output.append("vmovdqa ymm6, ymm8") |
| elif degree == 4: |
| output.append("vshufps zmm8, zmm4, zmm5, 214") |
| output.append("vpshufd zmm9, zmm4, 0x0F") |
| output.append("vpshufd zmm4, zmm8, 0x39") |
| output.append("vshufps zmm8, zmm6, zmm7, 250") |
| # k3 is initialized at the start of the kernel |
| output.append("vpblendmd zmm9 {k3}, zmm9, zmm8") |
| output.append("vpunpcklqdq zmm8, zmm7, zmm5") |
| # k4 is initialized at the start of the kernel |
| output.append("vpblendmd zmm8 {k4}, zmm8, zmm6") |
| output.append("vpshufd zmm8, zmm8, 0x78") |
| output.append("vpunpckhdq zmm5, zmm5, zmm7") |
| output.append("vpunpckldq zmm6, zmm6, zmm5") |
| output.append("vpshufd zmm7, zmm6, 0x1E") |
| output.append("vmovdqa32 zmm5, zmm9") |
| output.append("vmovdqa32 zmm6, zmm8") |
| else: |
| raise NotImplementedError |
| elif target.extension == AVX2: |
| assert degree == 2 |
| output.append("vshufps ymm8, ymm4, ymm5, 214") |
| output.append("vpshufd ymm9, ymm4, 0x0F") |
| output.append("vpshufd ymm4, ymm8, 0x39") |
| output.append("vshufps ymm8, ymm6, ymm7, 250") |
| output.append("vpblendd ymm9, ymm9, ymm8, 0xAA") |
| output.append("vpunpcklqdq ymm8, ymm7, ymm5") |
| output.append("vpblendd ymm8, ymm8, ymm6, 0x88") |
| output.append("vpshufd ymm8, ymm8, 0x78") |
| output.append("vpunpckhdq ymm5, ymm5, ymm7") |
| output.append("vpunpckldq ymm6, ymm6, ymm5") |
| output.append("vpshufd ymm7, ymm6, 0x1E") |
| output.append("vmovdqa ymm5, ymm9") |
| output.append("vmovdqa ymm6, ymm8") |
| elif target.extension == SSE41: |
| assert degree == 1 |
| output.append("movdqa xmm8, xmm4") |
| output.append("shufps xmm8, xmm5, 214") |
| output.append("pshufd xmm9, xmm4, 0x0F") |
| output.append("pshufd xmm4, xmm8, 0x39") |
| output.append("movdqa xmm8, xmm6") |
| output.append("shufps xmm8, xmm7, 250") |
| output.append("pblendw xmm9, xmm8, 0xCC") |
| output.append("movdqa xmm8, xmm7") |
| output.append("punpcklqdq xmm8, xmm5") |
| output.append("pblendw xmm8, xmm6, 0xC0") |
| output.append("pshufd xmm8, xmm8, 0x78") |
| output.append("punpckhdq xmm5, xmm7") |
| output.append("punpckldq xmm6, xmm5") |
| output.append("pshufd xmm7, xmm6, 0x1E") |
| output.append("movdqa xmm5, xmm9") |
| output.append("movdqa xmm6, xmm8") |
| elif target.extension == SSE2: |
| assert degree == 1 |
| output.append("movdqa xmm8, xmm4") |
| output.append("shufps xmm8, xmm5, 214") |
| output.append("pshufd xmm9, xmm4, 0x0F") |
| output.append("pshufd xmm4, xmm8, 0x39") |
| output.append("movdqa xmm8, xmm6") |
| output.append("shufps xmm8, xmm7, 250") |
| output.append("pand xmm9, xmmword ptr [PBLENDW_0x33_MASK+rip]") |
| output.append("pand xmm8, xmmword ptr [PBLENDW_0xCC_MASK+rip]") |
| output.append("por xmm9, xmm8") |
| output.append("movdqa xmm8, xmm7") |
| output.append("punpcklqdq xmm8, xmm5") |
| output.append("movdqa xmm10, xmm6") |
| output.append("pand xmm8, xmmword ptr [PBLENDW_0x3F_MASK+rip]") |
| output.append("pand xmm10, xmmword ptr [PBLENDW_0xC0_MASK+rip]") |
| output.append("por xmm8, xmm10") |
| output.append("pshufd xmm8, xmm8, 0x78") |
| output.append("punpckhdq xmm5, xmm7") |
| output.append("punpckldq xmm6, xmm5") |
| output.append("pshufd xmm7, xmm6, 0x1E") |
| output.append("movdqa xmm5, xmm9") |
| output.append("movdqa xmm6, xmm8") |
| else: |
| raise NotImplementedError |
| |
| |
| def kernel_2d_name(target, degree): |
| return f"blake3_{target.extension}_kernel_2d_{degree}" |
| |
| |
| # The two-dimensional kernel packs one or more *rows* of the state into a |
| # vector. For example, the AVX2 version of this kernel computes two inputs in |
| # parallel, and the caller needs to arrange their extended state words like |
| # this: |
| # |
| # ymm0: a0, a1, a2, a3, b0, b1, b2, b3 |
| # ymm1: a4, a5, a6, a7, b4, b5, b6, b7 |
| # ymm2: a8, a9, a10, a11, b8, b9, b10, b11 |
| # ymm3: a12, a13, a14, a15, b12, b13, b14, b15 |
| # |
| # In this arrangement, the rows need to be diagonalized and undiagonalized in |
| # each round. There's an important optimization for this, which applies to all |
| # ChaCha-derived functions. Intuitively, diagonalization in ChaCha looks like |
| # this: |
| # |
| # 0 1 2 3 ------------ no change -----------> 0 1 2 3 |
| # 4 5 6 7 ---- rotate one position left ----> 5 6 7 4 |
| # 8 9 10 11 ---- rotate two positions left ---> 10 11 8 9 |
| # 12 13 14 15 --- rotate three positions left --> 15 12 13 14 |
| # |
| # However, there's a performance benefit to doing it this way instead: |
| # |
| # 0 1 2 3 --- rotate three positions left --> 3 0 1 2 |
| # 4 5 6 7 ------------ no change -----------> 4 5 6 7 |
| # 8 9 10 11 ---- rotate one position left ----> 9 10 11 8 |
| # 12 13 14 15 ---- rotate two positions left ---> 14 15 12 13 |
| # |
| # That is, rather than keeping the first row fixed, keep the second row fixed. |
| # This gives the same columns, but in a different order. The second row is the |
| # last one touched by the G function, so leaving it unrotated saves latency. |
| # For more discussion, see: https://github.com/sneves/blake2-avx2/pull/4 |
| # |
| # The message words need to be arranged to match the state words. Again using |
| # AVX2 as an example, the caller needs to initially arrange the message words |
| # like this: |
| # |
| # ymm4: a0, a2, a4, a6, b0, b2, b4, b6 |
| # ymm5: a1, a3, a5, a7, b1, b3, b5, b7 |
| # ymm6: a14, a8, a10, a12, b14, b8, b10, b12 |
| # ymm7: a15, a9, a11, a13, b15, b9, b11, b13 |
| def kernel_2d(target, output, degree): |
| label = kernel_2d_name(target, degree) |
| output.append(f"{label}:") |
| # vpshufb indexes |
| if target.extension == SSE2: |
| assert degree == 1 |
| elif target.extension == SSE41: |
| assert degree == 1 |
| output.append(f"movaps xmm14, xmmword ptr [ROT8+rip]") |
| output.append(f"movaps xmm15, xmmword ptr [ROT16+rip]") |
| elif target.extension == AVX2: |
| assert degree == 2 |
| output.append(f"vbroadcasti128 ymm14, xmmword ptr [ROT16+rip]") |
| output.append(f"vbroadcasti128 ymm15, xmmword ptr [ROT8+rip]") |
| elif target.extension == AVX512: |
| assert degree in (1, 2, 4) |
| if degree == 4: |
| output.append(f"mov {target.scratch32(0)}, 43690") |
| output.append(f"kmovw k3, {target.scratch32(0)}") |
| output.append(f"mov {target.scratch32(0)}, 34952") |
| output.append(f"kmovw k4, {target.scratch32(0)}") |
| else: |
| raise NotImplementedError |
| for round_number in range(7): |
| if round_number > 0: |
| # Un-diagonalize and permute before each round except the first. |
| # compress_finish() will also partially un-diagonalize. |
| permute_message_rows(target, output, degree) |
| add_row(target, output, degree, dest=0, src=4) |
| add_row(target, output, degree, dest=0, src=1) |
| xor_row(target, output, degree, dest=3, src=0) |
| bitrotate_row(target, output, degree, reg=3, bits=16) |
| add_row(target, output, degree, dest=2, src=3) |
| xor_row(target, output, degree, dest=1, src=2) |
| bitrotate_row(target, output, degree, reg=1, bits=12) |
| add_row(target, output, degree, dest=0, src=5) |
| add_row(target, output, degree, dest=0, src=1) |
| xor_row(target, output, degree, dest=3, src=0) |
| bitrotate_row(target, output, degree, reg=3, bits=8) |
| add_row(target, output, degree, dest=2, src=3) |
| xor_row(target, output, degree, dest=1, src=2) |
| bitrotate_row(target, output, degree, reg=1, bits=7) |
| diagonalize_state_rows(target, output, degree) |
| add_row(target, output, degree, dest=0, src=6) |
| add_row(target, output, degree, dest=0, src=1) |
| xor_row(target, output, degree, dest=3, src=0) |
| bitrotate_row(target, output, degree, reg=3, bits=16) |
| add_row(target, output, degree, dest=2, src=3) |
| xor_row(target, output, degree, dest=1, src=2) |
| bitrotate_row(target, output, degree, reg=1, bits=12) |
| add_row(target, output, degree, dest=0, src=7) |
| add_row(target, output, degree, dest=0, src=1) |
| xor_row(target, output, degree, dest=3, src=0) |
| bitrotate_row(target, output, degree, reg=3, bits=8) |
| add_row(target, output, degree, dest=2, src=3) |
| xor_row(target, output, degree, dest=1, src=2) |
| bitrotate_row(target, output, degree, reg=1, bits=7) |
| undiagonalize_state_rows(target, output, degree) |
| # Xor the last two rows into the first two, but don't do the feed forward |
| # here. That's only done in the XOF case. |
| xor_row(target, output, degree, dest=0, src=2) |
| xor_row(target, output, degree, dest=1, src=3) |
| output.append(target.ret()) |
| |
| |
| def compress_setup(target, output): |
| if target.extension == AVX512: |
| # state words |
| output.append(f"vmovdqu xmm0, xmmword ptr [{target.arg64(0)}]") |
| output.append(f"vmovdqu xmm1, xmmword ptr [{target.arg64(0)}+0x10]") |
| # flags |
| output.append(f"shl {target.arg64(4)}, 32") |
| # block length |
| output.append(f"mov {target.arg32(3)}, {target.arg32(3)}") |
| output.append(f"or {target.arg64(3)}, {target.arg64(4)}") |
| # counter |
| output.append(f"vmovq xmm3, {target.arg64(2)}") |
| output.append(f"vmovq xmm4, {target.arg64(3)}") |
| output.append(f"vpunpcklqdq xmm3, xmm3, xmm4") |
| output.append(f"vmovaps xmm2, xmmword ptr [BLAKE3_IV+rip]") |
| # message words |
| # fmt: off |
| output.append(f"vmovups xmm8, xmmword ptr [{target.arg64(1)}]") # xmm8 = m0 m1 m2 m3 |
| output.append(f"vmovups xmm9, xmmword ptr [{target.arg64(1)}+0x10]") # xmm9 = m4 m5 m6 m7 |
| output.append(f"vshufps xmm4, xmm8, xmm9, 136") # xmm4 = m0 m2 m4 m6 |
| output.append(f"vshufps xmm5, xmm8, xmm9, 221") # xmm5 = m1 m3 m5 m7 |
| output.append(f"vmovups xmm8, xmmword ptr [{target.arg64(1)}+0x20]") # xmm8 = m8 m9 m10 m12 |
| output.append(f"vmovups xmm9, xmmword ptr [{target.arg64(1)}+0x30]") # xmm9 = m12 m13 m14 m15 |
| output.append(f"vshufps xmm6, xmm8, xmm9, 136") # xmm6 = m8 m10 m12 m14 |
| output.append(f"vshufps xmm7, xmm8, xmm9, 221") # xmm7 = m9 m11 m13 m15 |
| output.append(f"vpshufd xmm6, xmm6, 0x93") # xmm6 = m14 m8 m10 m12 |
| output.append(f"vpshufd xmm7, xmm7, 0x93") # xmm7 = m15 m9 m11 m13 |
| # fmt: on |
| elif target.extension in (SSE41, SSE2): |
| output.append(f"movups xmm0, xmmword ptr [{target.arg64(0)}]") |
| output.append(f"movups xmm1, xmmword ptr [{target.arg64(0)}+0x10]") |
| output.append(f"movaps xmm2, xmmword ptr [BLAKE3_IV+rip]") |
| output.append(f"shl {target.arg64(4)}, 32") |
| output.append(f"mov {target.arg32(3)}, {target.arg32(3)}") |
| output.append(f"or {target.arg64(3)}, {target.arg64(4)}") |
| output.append(f"vmovq xmm3, {target.arg64(2)}") |
| output.append(f"vmovq xmm4, {target.arg64(3)}") |
| output.append(f"punpcklqdq xmm3, xmm4") |
| output.append(f"movups xmm4, xmmword ptr [{target.arg64(1)}]") |
| output.append(f"movups xmm5, xmmword ptr [{target.arg64(1)}+0x10]") |
| output.append(f"movaps xmm8, xmm4") |
| output.append(f"shufps xmm4, xmm5, 136") |
| output.append(f"shufps xmm8, xmm5, 221") |
| output.append(f"movaps xmm5, xmm8") |
| output.append(f"movups xmm6, xmmword ptr [{target.arg64(1)}+0x20]") |
| output.append(f"movups xmm7, xmmword ptr [{target.arg64(1)}+0x30]") |
| output.append(f"movaps xmm8, xmm6") |
| output.append(f"shufps xmm6, xmm7, 136") |
| output.append(f"pshufd xmm6, xmm6, 0x93") |
| output.append(f"shufps xmm8, xmm7, 221") |
| output.append(f"pshufd xmm7, xmm8, 0x93") |
| else: |
| raise NotImplementedError |
| |
| |
| def compress_finish(target, output): |
| if target.extension == AVX512: |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(0)}], xmm0") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(0)}+0x10], xmm1") |
| elif target.extension in (SSE41, SSE2): |
| output.append(f"movups xmmword ptr [{target.arg64(0)}], xmm0") |
| output.append(f"movups xmmword ptr [{target.arg64(0)}+0x10], xmm1") |
| else: |
| raise NotImplementedError |
| |
| |
| def compress(target, output): |
| label = f"blake3_{target.extension}_compress" |
| output.append(f".global {label}") |
| output.append(f"{label}:") |
| compress_setup(target, output) |
| output.append(f"call {kernel_2d_name(target, 1)}") |
| compress_finish(target, output) |
| output.append(target.ret()) |
| |
| |
| def xof_setup_2d(target, output, degree): |
| if target.extension == AVX512: |
| if degree == 1: |
| # state words |
| output.append(f"vmovdqu xmm0, xmmword ptr [{target.arg64(0)}]") |
| output.append(f"vmovdqu xmm1, xmmword ptr [{target.arg64(0)}+0x10]") |
| # flags |
| output.append(f"shl {target.arg64(4)}, 32") |
| # block length |
| output.append(f"mov {target.arg32(3)}, {target.arg32(3)}") |
| output.append(f"or {target.arg64(3)}, {target.arg64(4)}") |
| # counter |
| output.append(f"vmovq xmm3, {target.arg64(2)}") |
| output.append(f"vmovq xmm4, {target.arg64(3)}") |
| output.append(f"vpunpcklqdq xmm3, xmm3, xmm4") |
| output.append(f"vmovaps xmm2, xmmword ptr [BLAKE3_IV+rip]") |
| # message words |
| # fmt: off |
| output.append(f"vmovups xmm8, xmmword ptr [{target.arg64(1)}]") # xmm8 = m0 m1 m2 m3 |
| output.append(f"vmovups xmm9, xmmword ptr [{target.arg64(1)}+0x10]") # xmm9 = m4 m5 m6 m7 |
| output.append(f"vshufps xmm4, xmm8, xmm9, 136") # xmm4 = m0 m2 m4 m6 |
| output.append(f"vshufps xmm5, xmm8, xmm9, 221") # xmm5 = m1 m3 m5 m7 |
| output.append(f"vmovups xmm8, xmmword ptr [{target.arg64(1)}+0x20]") # xmm8 = m8 m9 m10 m12 |
| output.append(f"vmovups xmm9, xmmword ptr [{target.arg64(1)}+0x30]") # xmm9 = m12 m13 m14 m15 |
| output.append(f"vshufps xmm6, xmm8, xmm9, 136") # xmm6 = m8 m10 m12 m14 |
| output.append(f"vshufps xmm7, xmm8, xmm9, 221") # xmm7 = m9 m11 m13 m15 |
| output.append(f"vpshufd xmm6, xmm6, 0x93") # xmm6 = m14 m8 m10 m12 |
| output.append(f"vpshufd xmm7, xmm7, 0x93") # xmm7 = m15 m9 m11 m13 |
| # fmt: on |
| elif degree == 2: |
| # Load the state words. |
| output.append(f"vbroadcasti128 ymm0, xmmword ptr [{target.arg64(0)}]") |
| output.append(f"vbroadcasti128 ymm1, xmmword ptr [{target.arg64(0)}+0x10]") |
| # Load the counter increments. |
| output.append(f"vmovdqa ymm4, ymmword ptr [INCREMENT_2D+rip]") |
| # Load the IV constants. |
| output.append(f"vbroadcasti128 ymm2, xmmword ptr [BLAKE3_IV+rip]") |
| # Broadcast the counter. |
| output.append(f"vpbroadcastq ymm5, {target.arg64(2)}") |
| # Add the counter increments to the counter. |
| output.append(f"vpaddq ymm6, ymm4, ymm5") |
| # Combine the block length and flags into a 64-bit word. |
| output.append(f"shl {target.arg64(4)}, 32") |
| output.append(f"mov {target.arg32(3)}, {target.arg32(3)}") |
| output.append(f"or {target.arg64(3)}, {target.arg64(4)}") |
| # Broadcast the block length and flags. |
| output.append(f"vpbroadcastq ymm7, {target.arg64(3)}") |
| # Blend the counter, block length, and flags. |
| output.append(f"vpblendd ymm3, ymm6, ymm7, 0xCC") |
| # Load and permute the message words. |
| output.append(f"vbroadcasti128 ymm8, xmmword ptr [{target.arg64(1)}]") |
| output.append(f"vbroadcasti128 ymm9, xmmword ptr [{target.arg64(1)}+0x10]") |
| output.append(f"vshufps ymm4, ymm8, ymm9, 136") |
| output.append(f"vshufps ymm5, ymm8, ymm9, 221") |
| output.append(f"vbroadcasti128 ymm8, xmmword ptr [{target.arg64(1)}+0x20]") |
| output.append(f"vbroadcasti128 ymm9, xmmword ptr [{target.arg64(1)}+0x30]") |
| output.append(f"vshufps ymm6, ymm8, ymm9, 136") |
| output.append(f"vshufps ymm7, ymm8, ymm9, 221") |
| output.append(f"vpshufd ymm6, ymm6, 0x93") |
| output.append(f"vpshufd ymm7, ymm7, 0x93") |
| elif degree == 4: |
| # Load the state words. |
| output.append(f"vbroadcasti32x4 zmm0, xmmword ptr [{target.arg64(0)}]") |
| output.append(f"vbroadcasti32x4 zmm1, xmmword ptr [{target.arg64(0)}+0x10]") |
| # Load the counter increments. |
| output.append(f"vmovdqa32 zmm4, zmmword ptr [INCREMENT_2D+rip]") |
| # Load the IV constants. |
| output.append(f"vbroadcasti32x4 zmm2, xmmword ptr [BLAKE3_IV+rip]") |
| # Broadcast the counter. |
| output.append(f"vpbroadcastq zmm5, {target.arg64(2)}") |
| # Add the counter increments to the counter. |
| output.append(f"vpaddq zmm6, zmm4, zmm5") |
| # Combine the block length and flags into a 64-bit word. |
| output.append(f"shl {target.arg64(4)}, 32") |
| output.append(f"mov {target.arg32(3)}, {target.arg32(3)}") |
| output.append(f"or {target.arg64(3)}, {target.arg64(4)}") |
| # Broadcast the block length and flags. |
| output.append(f"vpbroadcastq zmm7, {target.arg64(3)}") |
| # Blend the counter, block length, and flags. |
| output.append(f"mov {target.scratch32(0)}, 0xAA") |
| output.append(f"kmovw k2, {target.scratch32(0)}") |
| output.append(f"vpblendmq zmm3 {{k2}}, zmm6, zmm7") |
| # Load and permute the message words. |
| output.append(f"vbroadcasti32x4 zmm8, xmmword ptr [{target.arg64(1)}]") |
| output.append(f"vbroadcasti32x4 zmm9, xmmword ptr [{target.arg64(1)}+0x10]") |
| output.append(f"vshufps zmm4, zmm8, zmm9, 136") |
| output.append(f"vshufps zmm5, zmm8, zmm9, 221") |
| output.append(f"vbroadcasti32x4 zmm8, xmmword ptr [{target.arg64(1)}+0x20]") |
| output.append(f"vbroadcasti32x4 zmm9, xmmword ptr [{target.arg64(1)}+0x30]") |
| output.append(f"vshufps zmm6, zmm8, zmm9, 136") |
| output.append(f"vshufps zmm7, zmm8, zmm9, 221") |
| output.append(f"vpshufd zmm6, zmm6, 0x93") |
| output.append(f"vpshufd zmm7, zmm7, 0x93") |
| else: |
| raise NotImplementedError |
| elif target.extension == AVX2: |
| # Load the state words. |
| output.append(f"vbroadcasti128 ymm0, xmmword ptr [{target.arg64(0)}]") |
| output.append(f"vbroadcasti128 ymm1, xmmword ptr [{target.arg64(0)}+0x10]") |
| # Load the counter increments. |
| output.append(f"vmovdqa ymm4, ymmword ptr [INCREMENT_2D+rip]") |
| # Load the IV constants. |
| output.append(f"vbroadcasti128 ymm2, xmmword ptr [BLAKE3_IV+rip]") |
| # Broadcast the counter. |
| output.append(f"vpbroadcastq ymm5, {target.arg64(2)}") |
| # Add the counter increments to the counter. |
| output.append(f"vpaddq ymm6, ymm4, ymm5") |
| # Combine the block length and flags into a 64-bit word. |
| output.append(f"shl {target.arg64(4)}, 32") |
| output.append(f"mov {target.arg32(3)}, {target.arg32(3)}") |
| output.append(f"or {target.arg64(3)}, {target.arg64(4)}") |
| # Broadcast the block length and flags. |
| output.append(f"vpbroadcastq ymm7, {target.arg64(3)}") |
| # Blend the counter, block length, and flags. |
| output.append(f"vpblendd ymm3, ymm6, ymm7, 0xCC") |
| # Load and permute the message words. |
| output.append(f"vbroadcasti128 ymm8, xmmword ptr [{target.arg64(1)}]") |
| output.append(f"vbroadcasti128 ymm9, xmmword ptr [{target.arg64(1)}+0x10]") |
| output.append(f"vshufps ymm4, ymm8, ymm9, 136") |
| output.append(f"vshufps ymm5, ymm8, ymm9, 221") |
| output.append(f"vbroadcasti128 ymm8, xmmword ptr [{target.arg64(1)}+0x20]") |
| output.append(f"vbroadcasti128 ymm9, xmmword ptr [{target.arg64(1)}+0x30]") |
| output.append(f"vshufps ymm6, ymm8, ymm9, 136") |
| output.append(f"vshufps ymm7, ymm8, ymm9, 221") |
| output.append(f"vpshufd ymm6, ymm6, 0x93") |
| output.append(f"vpshufd ymm7, ymm7, 0x93") |
| elif target.extension in (SSE41, SSE2): |
| assert degree == 1 |
| output.append(f"movups xmm0, xmmword ptr [{target.arg64(0)}]") |
| output.append(f"movups xmm1, xmmword ptr [{target.arg64(0)}+0x10]") |
| output.append(f"movaps xmm2, xmmword ptr [BLAKE3_IV+rip]") |
| output.append(f"shl {target.arg64(4)}, 32") |
| output.append(f"mov {target.arg32(3)}, {target.arg32(3)}") |
| output.append(f"or {target.arg64(3)}, {target.arg64(4)}") |
| output.append(f"vmovq xmm3, {target.arg64(2)}") |
| output.append(f"vmovq xmm4, {target.arg64(3)}") |
| output.append(f"punpcklqdq xmm3, xmm4") |
| output.append(f"movups xmm4, xmmword ptr [{target.arg64(1)}]") |
| output.append(f"movups xmm5, xmmword ptr [{target.arg64(1)}+0x10]") |
| output.append(f"movaps xmm8, xmm4") |
| output.append(f"shufps xmm4, xmm5, 136") |
| output.append(f"shufps xmm8, xmm5, 221") |
| output.append(f"movaps xmm5, xmm8") |
| output.append(f"movups xmm6, xmmword ptr [{target.arg64(1)}+0x20]") |
| output.append(f"movups xmm7, xmmword ptr [{target.arg64(1)}+0x30]") |
| output.append(f"movaps xmm8, xmm6") |
| output.append(f"shufps xmm6, xmm7, 136") |
| output.append(f"pshufd xmm6, xmm6, 0x93") |
| output.append(f"shufps xmm8, xmm7, 221") |
| output.append(f"pshufd xmm7, xmm8, 0x93") |
| else: |
| raise NotImplementedError |
| |
| |
| def xof_stream_finish_2d(target, output, degree): |
| if target.extension == AVX512: |
| if degree == 1: |
| output.append(f"vpxor xmm2, xmm2, [{target.arg64(0)}]") |
| output.append(f"vpxor xmm3, xmm3, [{target.arg64(0)}+0x10]") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)}], xmm0") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)}+0x10], xmm1") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)}+0x20], xmm2") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)}+0x30], xmm3") |
| elif degree == 2: |
| output.append(f"vbroadcasti128 ymm4, xmmword ptr [{target.arg64(0)}]") |
| output.append(f"vpxor ymm2, ymm2, ymm4") |
| output.append(f"vbroadcasti128 ymm5, xmmword ptr [{target.arg64(0)} + 16]") |
| output.append(f"vpxor ymm3, ymm3, ymm5") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)} + 0 * 16], xmm0") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)} + 1 * 16], xmm1") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)} + 2 * 16], xmm2") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)} + 3 * 16], xmm3") |
| output.append(f"vextracti128 xmmword ptr [{target.arg64(5)}+4*16], ymm0, 1") |
| output.append(f"vextracti128 xmmword ptr [{target.arg64(5)}+5*16], ymm1, 1") |
| output.append(f"vextracti128 xmmword ptr [{target.arg64(5)}+6*16], ymm2, 1") |
| output.append(f"vextracti128 xmmword ptr [{target.arg64(5)}+7*16], ymm3, 1") |
| elif degree == 4: |
| output.append(f"vbroadcasti32x4 zmm4, xmmword ptr [{target.arg64(0)}]") |
| output.append(f"vpxord zmm2, zmm2, zmm4") |
| output.append(f"vbroadcasti32x4 zmm5, xmmword ptr [{target.arg64(0)} + 16]") |
| output.append(f"vpxord zmm3, zmm3, zmm5") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)} + 0 * 16], xmm0") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)} + 1 * 16], xmm1") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)} + 2 * 16], xmm2") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)} + 3 * 16], xmm3") |
| for i in range(1, 4): |
| for reg in range(0, 4): |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + {4*i+reg} * 16], zmm{reg}, {i}" |
| ) |
| else: |
| raise NotImplementedError |
| elif target.extension == AVX2: |
| output.append(f"vbroadcasti128 ymm4, xmmword ptr [{target.arg64(0)}]") |
| output.append(f"vpxor ymm2, ymm2, ymm4") |
| output.append(f"vbroadcasti128 ymm5, xmmword ptr [{target.arg64(0)} + 16]") |
| output.append(f"vpxor ymm3, ymm3, ymm5") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)} + 0 * 16], xmm0") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)} + 1 * 16], xmm1") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)} + 2 * 16], xmm2") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)} + 3 * 16], xmm3") |
| output.append(f"vextracti128 xmmword ptr [{target.arg64(5)} + 4 * 16], ymm0, 1") |
| output.append(f"vextracti128 xmmword ptr [{target.arg64(5)} + 5 * 16], ymm1, 1") |
| output.append(f"vextracti128 xmmword ptr [{target.arg64(5)} + 6 * 16], ymm2, 1") |
| output.append(f"vextracti128 xmmword ptr [{target.arg64(5)} + 7 * 16], ymm3, 1") |
| elif target.extension in (SSE41, SSE2): |
| assert degree == 1 |
| output.append(f"movdqu xmm4, xmmword ptr [{target.arg64(0)}]") |
| output.append(f"movdqu xmm5, xmmword ptr [{target.arg64(0)}+0x10]") |
| output.append(f"pxor xmm2, xmm4") |
| output.append(f"pxor xmm3, xmm5") |
| output.append(f"movups xmmword ptr [{target.arg64(5)}], xmm0") |
| output.append(f"movups xmmword ptr [{target.arg64(5)}+0x10], xmm1") |
| output.append(f"movups xmmword ptr [{target.arg64(5)}+0x20], xmm2") |
| output.append(f"movups xmmword ptr [{target.arg64(5)}+0x30], xmm3") |
| else: |
| raise NotImplementedError |
| |
| |
| def xof_xor_finish_2d(target, output, degree): |
| if target.extension == AVX512: |
| if degree == 1: |
| output.append(f"vpxor xmm2, xmm2, [{target.arg64(0)}]") |
| output.append(f"vpxor xmm3, xmm3, [{target.arg64(0)}+0x10]") |
| output.append(f"vpxor xmm0, xmm0, [{target.arg64(5)}]") |
| output.append(f"vpxor xmm1, xmm1, [{target.arg64(5)}+0x10]") |
| output.append(f"vpxor xmm2, xmm2, [{target.arg64(5)}+0x20]") |
| output.append(f"vpxor xmm3, xmm3, [{target.arg64(5)}+0x30]") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)}], xmm0") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)}+0x10], xmm1") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)}+0x20], xmm2") |
| output.append(f"vmovdqu xmmword ptr [{target.arg64(5)}+0x30], xmm3") |
| elif degree == 2: |
| output.append(f"vbroadcasti128 ymm4, xmmword ptr [{target.arg64(0)}]") |
| output.append(f"vpxor ymm2, ymm2, ymm4") |
| output.append(f"vbroadcasti128 ymm5, xmmword ptr [{target.arg64(0)} + 16]") |
| output.append(f"vpxor ymm3, ymm3, ymm5") |
| # Each vector now holds rows from two different states: |
| # ymm0: a0, a1, a2, a3, b0, b1, b2, b3 |
| # ymm1: a4, a5, a6, a7, b4, b5, b6, b7 |
| # ymm2: a8, a9, a10, a11, b8, b9, b10, b11 |
| # ymm3: a12, a13, a14, a15, b12, b13, b14, b15 |
| # We want to rearrange the 128-bit lanes like this, so we can load |
| # destination bytes and XOR them in directly. |
| # ymm4: a0, a1, a2, a3, a4, a5, a6, a7 |
| # ymm5: a8, a9, a10, a11, a12, a13, a14, a15 |
| # ymm6: b0, b1, b2, b3, b4, b5, b6, b7 |
| # ymm7: b8, b9, b10, b11, b12, b13, b14, b15 |
| output.append(f"vperm2f128 ymm4, ymm0, ymm1, {0b0010_0000}") # lower 128 |
| output.append(f"vperm2f128 ymm5, ymm2, ymm3, {0b0010_0000}") |
| output.append(f"vperm2f128 ymm6, ymm0, ymm1, {0b0011_0001}") # upper 128 |
| output.append(f"vperm2f128 ymm7, ymm2, ymm3, {0b0011_0001}") |
| # XOR in the bytes that are already in the destination. |
| output.append(f"vpxor ymm4, ymm4, ymmword ptr [{target.arg64(5)} + 0 * 32]") |
| output.append(f"vpxor ymm5, ymm5, ymmword ptr [{target.arg64(5)} + 1 * 32]") |
| output.append(f"vpxor ymm6, ymm6, ymmword ptr [{target.arg64(5)} + 2 * 32]") |
| output.append(f"vpxor ymm7, ymm7, ymmword ptr [{target.arg64(5)} + 3 * 32]") |
| # Write out the XOR results. |
| output.append(f"vmovdqu ymmword ptr [{target.arg64(5)} + 0 * 32], ymm4") |
| output.append(f"vmovdqu ymmword ptr [{target.arg64(5)} + 1 * 32], ymm5") |
| output.append(f"vmovdqu ymmword ptr [{target.arg64(5)} + 2 * 32], ymm6") |
| output.append(f"vmovdqu ymmword ptr [{target.arg64(5)} + 3 * 32], ymm7") |
| elif degree == 4: |
| output.append(f"vbroadcasti32x4 zmm4, xmmword ptr [{target.arg64(0)}]") |
| output.append(f"vpxord zmm2, zmm2, zmm4") |
| output.append(f"vbroadcasti32x4 zmm5, xmmword ptr [{target.arg64(0)} + 16]") |
| output.append(f"vpxord zmm3, zmm3, zmm5") |
| # Each vector now holds rows from four different states: |
| # zmm0: a0, a1, a2, a3, b0, b1, b2, b3, c0, c1, c2, c3, d0, d1, d2, d3 |
| # zmm1: a4, a5, a6, a7, b4, b5, b6, b7, c4, c5, c6, c7, d4, d5, d6, d7 |
| # zmm2: a8, a9, a10, a11, b8, b9, b10, b11, c8, c9, c10, c11, d8, d9, d10, d11 |
| # zmm3: a12, a13, a14, a15, b12, b13, b14, b15, c12, c13, c14, c15, d12, d13, d14, d15 |
| # We want to rearrange the 128-bit lanes like this, so we can load |
| # destination bytes and XOR them in directly. |
| # zmm0: a0, a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11, a12, a13, a14, a15 |
| # zmm1: b0, b1, b2, b3, b4, b5, b6, b7, b8, b9, b10, b11, b12, b13, b14, b15 |
| # zmm2: c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15 |
| # zmm3: d0, d1, d2, d3, d4, d5, d6, d7, d8, d9, d10, d11, d12, d13, d14, d15 |
| # |
| # This first interleaving of 256-bit lanes produces vectors like: |
| # zmm4: a0, a1, a2, a3, b0, b1, b2, b3, a4, a5, a6, a7, b4, b5, b6, b7 |
| output.append(f"vshufi32x4 zmm4, zmm0, zmm1, {0b0100_0100}") # low 256 |
| output.append(f"vshufi32x4 zmm5, zmm0, zmm1, {0b1110_1110}") # high 256 |
| output.append(f"vshufi32x4 zmm6, zmm2, zmm3, {0b0100_0100}") |
| output.append(f"vshufi32x4 zmm7, zmm2, zmm3, {0b1110_1110}") |
| # And this second interleaving of 128-bit lanes within each 256-bit |
| # lane produces the vectors we want. |
| output.append(f"vshufi32x4 zmm0, zmm4, zmm6, {0b1000_1000}") # low 128 |
| output.append(f"vshufi32x4 zmm1, zmm4, zmm6, {0b1101_1101}") # high 128 |
| output.append(f"vshufi32x4 zmm2, zmm5, zmm7, {0b1000_1000}") |
| output.append(f"vshufi32x4 zmm3, zmm5, zmm7, {0b1101_1101}") |
| # XOR in the bytes that are already in the destination. |
| output.append(f"vpxord zmm0, zmm0, zmmword ptr [{target.arg64(5)} + 0*64]") |
| output.append(f"vpxord zmm1, zmm1, zmmword ptr [{target.arg64(5)} + 1*64]") |
| output.append(f"vpxord zmm2, zmm2, zmmword ptr [{target.arg64(5)} + 2*64]") |
| output.append(f"vpxord zmm3, zmm3, zmmword ptr [{target.arg64(5)} + 3*64]") |
| # Write out the XOR results. |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 0*64], zmm0") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 1*64], zmm1") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 2*64], zmm2") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 3*64], zmm3") |
| else: |
| raise NotImplementedError |
| elif target.extension == AVX2: |
| output.append(f"vbroadcasti128 ymm4, xmmword ptr [{target.arg64(0)}]") |
| output.append(f"vpxor ymm2, ymm2, ymm4") |
| output.append(f"vbroadcasti128 ymm5, xmmword ptr [{target.arg64(0)} + 16]") |
| output.append(f"vpxor ymm3, ymm3, ymm5") |
| # Each vector now holds rows from two different states: |
| # ymm0: a0, a1, a2, a3, b0, b1, b2, b3 |
| # ymm1: a4, a5, a6, a7, b4, b5, b6, b7 |
| # ymm2: a8, a9, a10, a11, b8, b9, b10, b11 |
| # ymm3: a12, a13, a14, a15, b12, b13, b14, b15 |
| # We want to rearrange the 128-bit lanes like this, so we can load |
| # destination bytes and XOR them in directly. |
| # ymm4: a0, a1, a2, a3, a4, a5, a6, a7 |
| # ymm5: a8, a9, a10, a11, a12, a13, a14, a15 |
| # ymm6: b0, b1, b2, b3, b4, b5, b6, b7 |
| # ymm7: b8, b9, b10, b11, b12, b13, b14, b15 |
| output.append(f"vperm2f128 ymm4, ymm0, ymm1, {0b0010_0000}") # lower 128 |
| output.append(f"vperm2f128 ymm5, ymm2, ymm3, {0b0010_0000}") |
| output.append(f"vperm2f128 ymm6, ymm0, ymm1, {0b0011_0001}") # upper 128 |
| output.append(f"vperm2f128 ymm7, ymm2, ymm3, {0b0011_0001}") |
| # XOR in the bytes that are already in the destination. |
| output.append(f"vpxor ymm4, ymm4, ymmword ptr [{target.arg64(5)} + 0 * 32]") |
| output.append(f"vpxor ymm5, ymm5, ymmword ptr [{target.arg64(5)} + 1 * 32]") |
| output.append(f"vpxor ymm6, ymm6, ymmword ptr [{target.arg64(5)} + 2 * 32]") |
| output.append(f"vpxor ymm7, ymm7, ymmword ptr [{target.arg64(5)} + 3 * 32]") |
| # Write out the XOR results. |
| output.append(f"vmovdqu ymmword ptr [{target.arg64(5)} + 0 * 32], ymm4") |
| output.append(f"vmovdqu ymmword ptr [{target.arg64(5)} + 1 * 32], ymm5") |
| output.append(f"vmovdqu ymmword ptr [{target.arg64(5)} + 2 * 32], ymm6") |
| output.append(f"vmovdqu ymmword ptr [{target.arg64(5)} + 3 * 32], ymm7") |
| elif target.extension in (SSE41, SSE2): |
| assert degree == 1 |
| output.append(f"movdqu xmm4, xmmword ptr [{target.arg64(0)}]") |
| output.append(f"movdqu xmm5, xmmword ptr [{target.arg64(0)}+0x10]") |
| output.append(f"pxor xmm2, xmm4") |
| output.append(f"pxor xmm3, xmm5") |
| output.append(f"movdqu xmm4, [{target.arg64(5)}]") |
| output.append(f"movdqu xmm5, [{target.arg64(5)}+0x10]") |
| output.append(f"movdqu xmm6, [{target.arg64(5)}+0x20]") |
| output.append(f"movdqu xmm7, [{target.arg64(5)}+0x30]") |
| output.append(f"pxor xmm0, xmm4") |
| output.append(f"pxor xmm1, xmm5") |
| output.append(f"pxor xmm2, xmm6") |
| output.append(f"pxor xmm3, xmm7") |
| output.append(f"movups xmmword ptr [{target.arg64(5)}], xmm0") |
| output.append(f"movups xmmword ptr [{target.arg64(5)}+0x10], xmm1") |
| output.append(f"movups xmmword ptr [{target.arg64(5)}+0x20], xmm2") |
| output.append(f"movups xmmword ptr [{target.arg64(5)}+0x30], xmm3") |
| else: |
| raise NotImplementedError |
| |
| |
| def g_function_3d(target, output, degree, columns, msg_words1, msg_words2): |
| if target.extension == SSE41: |
| assert degree == 4 |
| elif target.extension == AVX2: |
| assert degree == 8 |
| elif target.extension == AVX512: |
| assert degree in (8, 16) |
| else: |
| raise NotImplementedError |
| for (column, m1) in zip(columns, msg_words1): |
| add_row(target, output, degree, dest=column[0], src=m1) |
| for column in columns: |
| add_row(target, output, degree, dest=column[0], src=column[1]) |
| for column in columns: |
| xor_row(target, output, degree, dest=column[3], src=column[0]) |
| for column in columns: |
| bitrotate_row(target, output, degree, reg=column[3], bits=16) |
| for column in columns: |
| add_row(target, output, degree, dest=column[2], src=column[3]) |
| for column in columns: |
| xor_row(target, output, degree, dest=column[1], src=column[2]) |
| for column in columns: |
| bitrotate_row(target, output, degree, reg=column[1], bits=12) |
| for (column, m2) in zip(columns, msg_words2): |
| add_row(target, output, degree, dest=column[0], src=m2) |
| for column in columns: |
| add_row(target, output, degree, dest=column[0], src=column[1]) |
| for column in columns: |
| xor_row(target, output, degree, dest=column[3], src=column[0]) |
| for column in columns: |
| bitrotate_row(target, output, degree, reg=column[3], bits=8) |
| for column in columns: |
| add_row(target, output, degree, dest=column[2], src=column[3]) |
| for column in columns: |
| xor_row(target, output, degree, dest=column[1], src=column[2]) |
| for column in columns: |
| bitrotate_row(target, output, degree, reg=column[1], bits=7) |
| |
| |
| def kernel_3d_name(target, degree): |
| return f"blake3_{target.extension}_kernel_3d_{degree}" |
| |
| |
| def kernel_3d(target, output, degree): |
| label = kernel_3d_name(target, degree) |
| output.append(f"{label}:") |
| if target.extension == SSE41: |
| assert degree == 4 |
| elif target.extension == AVX2: |
| assert degree == 8 |
| elif target.extension == AVX512: |
| assert degree in (8, 16) |
| else: |
| raise NotImplementedError |
| for round_number in range(7): |
| straight_columns = [ |
| [0, 4, 8, 12], |
| [1, 5, 9, 13], |
| [2, 6, 10, 14], |
| [3, 7, 11, 15], |
| ] |
| msg_words1 = [16 + MESSAGE_SCHEDULE[round_number][i] for i in [0, 2, 4, 6]] |
| msg_words2 = [16 + MESSAGE_SCHEDULE[round_number][i] for i in [1, 3, 5, 7]] |
| g_function_3d(target, output, degree, straight_columns, msg_words1, msg_words2) |
| diagonal_columns = [ |
| [0, 5, 10, 15], |
| [1, 6, 11, 12], |
| [2, 7, 8, 13], |
| [3, 4, 9, 14], |
| ] |
| msg_words1 = [16 + MESSAGE_SCHEDULE[round_number][i] for i in [8, 10, 12, 14]] |
| msg_words2 = [16 + MESSAGE_SCHEDULE[round_number][i] for i in [9, 11, 13, 15]] |
| g_function_3d(target, output, degree, diagonal_columns, msg_words1, msg_words2) |
| # Xor the last two rows into the first two, but don't do the feed forward |
| # here. That's only done in the XOF case. |
| for dest in range(8): |
| xor_row(target, output, degree, dest=dest, src=dest + 8) |
| output.append(target.ret()) |
| |
| |
| def xof_setup_3d(target, output, degree): |
| if target.extension == AVX512: |
| if degree == 16: |
| # Load vpermi2d indexes into the counter registers. |
| output.append(f"vmovdqa32 zmm12, zmmword ptr [rip + EVEN_INDEXES]") |
| output.append(f"vmovdqa32 zmm13, zmmword ptr [rip + ODD_INDEXES]") |
| # Load the state words. |
| for i in range(8): |
| output.append( |
| f"vpbroadcastd zmm{i}, dword ptr [{target.arg64(0)}+{4*i}]" |
| ) |
| # Load the message words. |
| for i in range(16): |
| output.append( |
| f"vpbroadcastd zmm{i+16}, dword ptr [{target.arg64(1)}+{4*i}]" |
| ) |
| # Load the 64-bit counter increments into a temporary register. |
| output.append(f"vmovdqa64 zmm8, zmmword ptr [INCREMENT_3D+rip]") |
| # Broadcast the counter and add it to the increments. This gives |
| # the first 8 counter values. |
| output.append(f"vpbroadcastq zmm9, {target.arg64(2)}") |
| output.append(f"vpaddq zmm9, zmm9, zmm8") |
| # Increment the counter and repeat that for the last 8 counter values. |
| output.append(f"add {target.arg64(2)}, 8") |
| output.append(f"vpbroadcastq zmm10, {target.arg64(2)}") |
| output.append(f"vpaddq zmm10, zmm10, zmm8") |
| # Extract the lower and upper halves of the counter words, using |
| # the permutation tables loaded above. |
| output.append(f"vpermi2d zmm12, zmm9, zmm10") |
| output.append(f"vpermi2d zmm13, zmm9, zmm10") |
| # Load the block length. |
| output.append(f"vpbroadcastd zmm14, {target.arg32(3)}") |
| # Load the domain flags. |
| output.append(f"vpbroadcastd zmm15, {target.arg32(4)}") |
| # Load the IV constants. |
| for i in range(4): |
| output.append(f"vpbroadcastd zmm{i+8}, dword ptr [BLAKE3_IV+rip+{4*i}]") |
| else: |
| raise NotImplementedError |
| else: |
| raise NotImplementedError |
| |
| |
| def xof_stream_finish_3d(target, output, degree): |
| if target.extension == AVX512: |
| if degree == 16: |
| # Re-broadcast the input CV and feed it forward into the second half of the state. |
| output.append(f"vpbroadcastd zmm16, dword ptr [{target.arg64(0)} + 0 * 4]") |
| output.append(f"vpxord zmm8, zmm8, zmm16") |
| output.append(f"vpbroadcastd zmm17, dword ptr [{target.arg64(0)} + 1 * 4]") |
| output.append(f"vpxord zmm9, zmm9, zmm17") |
| output.append(f"vpbroadcastd zmm18, dword ptr [{target.arg64(0)} + 2 * 4]") |
| output.append(f"vpxord zmm10, zmm10, zmm18") |
| output.append(f"vpbroadcastd zmm19, dword ptr [{target.arg64(0)} + 3 * 4]") |
| output.append(f"vpxord zmm11, zmm11, zmm19") |
| output.append(f"vpbroadcastd zmm20, dword ptr [{target.arg64(0)} + 4 * 4]") |
| output.append(f"vpxord zmm12, zmm12, zmm20") |
| output.append(f"vpbroadcastd zmm21, dword ptr [{target.arg64(0)} + 5 * 4]") |
| output.append(f"vpxord zmm13, zmm13, zmm21") |
| output.append(f"vpbroadcastd zmm22, dword ptr [{target.arg64(0)} + 6 * 4]") |
| output.append(f"vpxord zmm14, zmm14, zmm22") |
| output.append(f"vpbroadcastd zmm23, dword ptr [{target.arg64(0)} + 7 * 4]") |
| output.append(f"vpxord zmm15, zmm15, zmm23") |
| # zmm0-zmm15 now contain the final extended state vectors, transposed. We need to un-transpose |
| # them before we write them out. As with blake3_avx512_blocks_16, we prefer to avoid expensive |
| # operations across 128-bit lanes, so we do a couple of interleaving passes and then write out |
| # 128 bits at a time. |
| # |
| # First, interleave 32-bit words. Use zmm16-zmm31 to hold the intermediate results. This |
| # takes the input vectors like: |
| # |
| # a0, b0, c0, d0, e0, f0, g0, h0, i0, j0, k0, l0, m0, n0, o0, p0 |
| # |
| # And produces vectors like: |
| # |
| # a0, a1, b0, b1, e0, e1, g0, g1, i0, i1, k0, k1, m0, m1, o0, o1 |
| # |
| # Then interleave 64-bit words back into zmm0-zmm15, producing vectors like: |
| # |
| # a0, a1, a2, a3, e0, e1, e2, e3, i0, i1, i2, i3, m0, m1, m2, m3 |
| # |
| # Finally, write out each 128-bit group, unaligned. |
| output.append(f"vpunpckldq zmm16, zmm0, zmm1") |
| output.append(f"vpunpckhdq zmm17, zmm0, zmm1") |
| output.append(f"vpunpckldq zmm18, zmm2, zmm3") |
| output.append(f"vpunpckhdq zmm19, zmm2, zmm3") |
| output.append(f"vpunpcklqdq zmm0, zmm16, zmm18") |
| output.append(f"vmovdqu32 xmmword ptr [{target.arg64(5)} + 0 * 16], xmm0") |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 16 * 16], zmm0, 1" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 32 * 16], zmm0, 2" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 48 * 16], zmm0, 3" |
| ) |
| output.append(f"vpunpckhqdq zmm1, zmm16, zmm18") |
| output.append(f"vmovdqu32 xmmword ptr [{target.arg64(5)} + 4 * 16], xmm1") |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 20 * 16], zmm1, 1" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 36 * 16], zmm1, 2" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 52 * 16], zmm1, 3" |
| ) |
| output.append(f"vpunpcklqdq zmm2, zmm17, zmm19") |
| output.append(f"vmovdqu32 xmmword ptr [{target.arg64(5)} + 8 * 16], xmm2") |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 24 * 16], zmm2, 1" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 40 * 16], zmm2, 2" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 56 * 16], zmm2, 3" |
| ) |
| output.append(f"vpunpckhqdq zmm3, zmm17, zmm19") |
| output.append(f"vmovdqu32 xmmword ptr [{target.arg64(5)} + 12 * 16], xmm3") |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 28 * 16], zmm3, 1" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 44 * 16], zmm3, 2" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 60 * 16], zmm3, 3" |
| ) |
| output.append(f"vpunpckldq zmm20, zmm4, zmm5") |
| output.append(f"vpunpckhdq zmm21, zmm4, zmm5") |
| output.append(f"vpunpckldq zmm22, zmm6, zmm7") |
| output.append(f"vpunpckhdq zmm23, zmm6, zmm7") |
| output.append(f"vpunpcklqdq zmm4, zmm20, zmm22") |
| output.append(f"vmovdqu32 xmmword ptr [{target.arg64(5)} + 1 * 16], xmm4") |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 17 * 16], zmm4, 1" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 33 * 16], zmm4, 2" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 49 * 16], zmm4, 3" |
| ) |
| output.append(f"vpunpckhqdq zmm5, zmm20, zmm22") |
| output.append(f"vmovdqu32 xmmword ptr [{target.arg64(5)} + 5 * 16], xmm5") |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 21 * 16], zmm5, 1" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 37 * 16], zmm5, 2" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 53 * 16], zmm5, 3" |
| ) |
| output.append(f"vpunpcklqdq zmm6, zmm21, zmm23") |
| output.append(f"vmovdqu32 xmmword ptr [{target.arg64(5)} + 9 * 16], xmm6") |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 25 * 16], zmm6, 1" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 41 * 16], zmm6, 2" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 57 * 16], zmm6, 3" |
| ) |
| output.append(f"vpunpckhqdq zmm7, zmm21, zmm23") |
| output.append(f"vmovdqu32 xmmword ptr [{target.arg64(5)} + 13 * 16], xmm7") |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 29 * 16], zmm7, 1" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 45 * 16], zmm7, 2" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 61 * 16], zmm7, 3" |
| ) |
| output.append(f"vpunpckldq zmm24, zmm8, zmm9") |
| output.append(f"vpunpckhdq zmm25, zmm8, zmm9") |
| output.append(f"vpunpckldq zmm26, zmm10, zmm11") |
| output.append(f"vpunpckhdq zmm27, zmm10, zmm11") |
| output.append(f"vpunpcklqdq zmm8, zmm24, zmm26") |
| output.append(f"vmovdqu32 xmmword ptr [{target.arg64(5)} + 2 * 16], xmm8") |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 18 * 16], zmm8, 1" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 34 * 16], zmm8, 2" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 50 * 16], zmm8, 3" |
| ) |
| output.append(f"vpunpckhqdq zmm9, zmm24, zmm26") |
| output.append(f"vmovdqu32 xmmword ptr [{target.arg64(5)} + 6 * 16], xmm9") |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 22 * 16], zmm9, 1" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 38 * 16], zmm9, 2" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 54 * 16], zmm9, 3" |
| ) |
| output.append(f"vpunpcklqdq zmm10, zmm25, zmm27") |
| output.append(f"vmovdqu32 xmmword ptr [{target.arg64(5)} + 10 * 16], xmm10") |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 26 * 16], zmm10, 1" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 42 * 16], zmm10, 2" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 58 * 16], zmm10, 3" |
| ) |
| output.append(f"vpunpckhqdq zmm11, zmm25, zmm27") |
| output.append(f"vmovdqu32 xmmword ptr [{target.arg64(5)} + 14 * 16], xmm11") |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 30 * 16], zmm11, 1" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 46 * 16], zmm11, 2" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 62 * 16], zmm11, 3" |
| ) |
| output.append(f"vpunpckldq zmm28, zmm12, zmm13") |
| output.append(f"vpunpckhdq zmm29, zmm12, zmm13") |
| output.append(f"vpunpckldq zmm30, zmm14, zmm15") |
| output.append(f"vpunpckhdq zmm31, zmm14, zmm15") |
| output.append(f"vpunpcklqdq zmm12, zmm28, zmm30") |
| output.append(f"vmovdqu32 xmmword ptr [{target.arg64(5)} + 3 * 16], xmm12") |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 19 * 16], zmm12, 1" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 35 * 16], zmm12, 2" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 51 * 16], zmm12, 3" |
| ) |
| output.append(f"vpunpckhqdq zmm13, zmm28, zmm30") |
| output.append(f"vmovdqu32 xmmword ptr [{target.arg64(5)} + 7 * 16], xmm13") |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 23 * 16], zmm13, 1" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 39 * 16], zmm13, 2" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 55 * 16], zmm13, 3" |
| ) |
| output.append(f"vpunpcklqdq zmm14, zmm29, zmm31") |
| output.append(f"vmovdqu32 xmmword ptr [{target.arg64(5)} + 11 * 16], xmm14") |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 27 * 16], zmm14, 1" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 43 * 16], zmm14, 2" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 59 * 16], zmm14, 3" |
| ) |
| output.append(f"vpunpckhqdq zmm15, zmm29, zmm31") |
| output.append(f"vmovdqu32 xmmword ptr [{target.arg64(5)} + 15 * 16], xmm15") |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 31 * 16], zmm15, 1" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 47 * 16], zmm15, 2" |
| ) |
| output.append( |
| f"vextracti32x4 xmmword ptr [{target.arg64(5)} + 63 * 16], zmm15, 3" |
| ) |
| else: |
| raise NotImplementedError |
| else: |
| raise NotImplementedError |
| |
| |
| def xof_xor_finish_3d(target, output, degree): |
| if target.extension == AVX512: |
| if degree == 16: |
| # Re-broadcast the input CV and feed it forward into the second half of the state. |
| output.append(f"vpbroadcastd zmm16, dword ptr [{target.arg64(0)} + 0 * 4]") |
| output.append(f"vpxord zmm8, zmm8, zmm16") |
| output.append(f"vpbroadcastd zmm17, dword ptr [{target.arg64(0)} + 1 * 4]") |
| output.append(f"vpxord zmm9, zmm9, zmm17") |
| output.append(f"vpbroadcastd zmm18, dword ptr [{target.arg64(0)} + 2 * 4]") |
| output.append(f"vpxord zmm10, zmm10, zmm18") |
| output.append(f"vpbroadcastd zmm19, dword ptr [{target.arg64(0)} + 3 * 4]") |
| output.append(f"vpxord zmm11, zmm11, zmm19") |
| output.append(f"vpbroadcastd zmm20, dword ptr [{target.arg64(0)} + 4 * 4]") |
| output.append(f"vpxord zmm12, zmm12, zmm20") |
| output.append(f"vpbroadcastd zmm21, dword ptr [{target.arg64(0)} + 5 * 4]") |
| output.append(f"vpxord zmm13, zmm13, zmm21") |
| output.append(f"vpbroadcastd zmm22, dword ptr [{target.arg64(0)} + 6 * 4]") |
| output.append(f"vpxord zmm14, zmm14, zmm22") |
| output.append(f"vpbroadcastd zmm23, dword ptr [{target.arg64(0)} + 7 * 4]") |
| output.append(f"vpxord zmm15, zmm15, zmm23") |
| # zmm0-zmm15 now contain the final extended state vectors, transposed. We need to un-transpose |
| # them before we write them out. Unlike blake3_avx512_xof_stream_16, we do a complete |
| # un-transpose here, to make the xor step easier. |
| # |
| # First interleave 32-bit words. This takes vectors like: |
| # |
| # a0, b0, c0, d0, e0, f0, g0, h0, i0, j0, k0, l0, m0, n0, o0, p0 |
| # |
| # And produces vectors like: |
| # |
| # a0, a1, b0, b1, e0, e1, g0, g1, i0, i1, k0, k1, m0, m1, o0, o1 |
| output.append(f"vpunpckldq zmm16, zmm0, zmm1") |
| output.append(f"vpunpckhdq zmm17, zmm0, zmm1") |
| output.append(f"vpunpckldq zmm18, zmm2, zmm3") |
| output.append(f"vpunpckhdq zmm19, zmm2, zmm3") |
| output.append(f"vpunpckldq zmm20, zmm4, zmm5") |
| output.append(f"vpunpckhdq zmm21, zmm4, zmm5") |
| output.append(f"vpunpckldq zmm22, zmm6, zmm7") |
| output.append(f"vpunpckhdq zmm23, zmm6, zmm7") |
| output.append(f"vpunpckldq zmm24, zmm8, zmm9") |
| output.append(f"vpunpckhdq zmm25, zmm8, zmm9") |
| output.append(f"vpunpckldq zmm26, zmm10, zmm11") |
| output.append(f"vpunpckhdq zmm27, zmm10, zmm11") |
| output.append(f"vpunpckldq zmm28, zmm12, zmm13") |
| output.append(f"vpunpckhdq zmm29, zmm12, zmm13") |
| output.append(f"vpunpckldq zmm30, zmm14, zmm15") |
| output.append(f"vpunpckhdq zmm31, zmm14, zmm15") |
| # Then interleave 64-bit words, producing vectors like: |
| # |
| # a0, a1, a2, a3, e0, e1, e2, e3, i0, i1, i2, i3, m0, m1, m2, m3 |
| output.append(f"vpunpcklqdq zmm0, zmm16, zmm18") |
| output.append(f"vpunpckhqdq zmm1, zmm16, zmm18") |
| output.append(f"vpunpcklqdq zmm2, zmm17, zmm19") |
| output.append(f"vpunpckhqdq zmm3, zmm17, zmm19") |
| output.append(f"vpunpcklqdq zmm4, zmm20, zmm22") |
| output.append(f"vpunpckhqdq zmm5, zmm20, zmm22") |
| output.append(f"vpunpcklqdq zmm6, zmm21, zmm23") |
| output.append(f"vpunpckhqdq zmm7, zmm21, zmm23") |
| output.append(f"vpunpcklqdq zmm8, zmm24, zmm26") |
| output.append(f"vpunpckhqdq zmm9, zmm24, zmm26") |
| output.append(f"vpunpcklqdq zmm10, zmm25, zmm27") |
| output.append(f"vpunpckhqdq zmm11, zmm25, zmm27") |
| output.append(f"vpunpcklqdq zmm12, zmm28, zmm30") |
| output.append(f"vpunpckhqdq zmm13, zmm28, zmm30") |
| output.append(f"vpunpcklqdq zmm14, zmm29, zmm31") |
| output.append(f"vpunpckhqdq zmm15, zmm29, zmm31") |
| # Then interleave 128-bit lanes, producing vectors like: |
| # |
| # a0, a1, a2, a3, i0, i1, i2, i3, a4, a5, a6, a7, i4, i5, i6, i7 |
| output.append( |
| "vshufi32x4 zmm16, zmm0, zmm4, 0x88" |
| ) # lo lanes: 0x88 = 0b10001000 = (0, 2, 0, 2) |
| output.append(f"vshufi32x4 zmm17, zmm1, zmm5, 0x88") |
| output.append(f"vshufi32x4 zmm18, zmm2, zmm6, 0x88") |
| output.append(f"vshufi32x4 zmm19, zmm3, zmm7, 0x88") |
| output.append( |
| "vshufi32x4 zmm20, zmm0, zmm4, 0xdd" |
| ) # hi lanes: 0xdd = 0b11011101 = (1, 3, 1, 3) |
| output.append(f"vshufi32x4 zmm21, zmm1, zmm5, 0xdd") |
| output.append(f"vshufi32x4 zmm22, zmm2, zmm6, 0xdd") |
| output.append(f"vshufi32x4 zmm23, zmm3, zmm7, 0xdd") |
| output.append(f"vshufi32x4 zmm24, zmm8, zmm12, 0x88") # lo lanes |
| output.append(f"vshufi32x4 zmm25, zmm9, zmm13, 0x88") |
| output.append(f"vshufi32x4 zmm26, zmm10, zmm14, 0x88") |
| output.append(f"vshufi32x4 zmm27, zmm11, zmm15, 0x88") |
| output.append(f"vshufi32x4 zmm28, zmm8, zmm12, 0xdd") # hi lanes |
| output.append(f"vshufi32x4 zmm29, zmm9, zmm13, 0xdd") |
| output.append(f"vshufi32x4 zmm30, zmm10, zmm14, 0xdd") |
| output.append(f"vshufi32x4 zmm31, zmm11, zmm15, 0xdd") |
| # Finally interleave 128-bit lanes again (the same permutation as the previous pass, but |
| # different inputs), producing vectors like: |
| # |
| # a0, a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11, a12, a13, a14, a15 |
| output.append(f"vshufi32x4 zmm0, zmm16, zmm24, 0x88") # lo lanes |
| output.append(f"vshufi32x4 zmm1, zmm17, zmm25, 0x88") |
| output.append(f"vshufi32x4 zmm2, zmm18, zmm26, 0x88") |
| output.append(f"vshufi32x4 zmm3, zmm19, zmm27, 0x88") |
| output.append(f"vshufi32x4 zmm4, zmm20, zmm28, 0x88") |
| output.append(f"vshufi32x4 zmm5, zmm21, zmm29, 0x88") |
| output.append(f"vshufi32x4 zmm6, zmm22, zmm30, 0x88") |
| output.append(f"vshufi32x4 zmm7, zmm23, zmm31, 0x88") |
| output.append(f"vshufi32x4 zmm8, zmm16, zmm24, 0xdd") # hi lanes |
| output.append(f"vshufi32x4 zmm9, zmm17, zmm25, 0xdd") |
| output.append(f"vshufi32x4 zmm10, zmm18, zmm26, 0xdd") |
| output.append(f"vshufi32x4 zmm11, zmm19, zmm27, 0xdd") |
| output.append(f"vshufi32x4 zmm12, zmm20, zmm28, 0xdd") |
| output.append(f"vshufi32x4 zmm13, zmm21, zmm29, 0xdd") |
| output.append(f"vshufi32x4 zmm14, zmm22, zmm30, 0xdd") |
| output.append(f"vshufi32x4 zmm15, zmm23, zmm31, 0xdd") |
| # zmm0-zmm15 now contain the fully un-transposed state words. Load each 64 block on input |
| # (unaligned), perform the xor, and write out the result (again unaligned). |
| output.append(f"vmovdqu32 zmm16, zmmword ptr [{target.arg64(5)} + 0 * 64]") |
| output.append(f"vpxord zmm0, zmm0, zmm16") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 0 * 64], zmm0") |
| output.append(f"vmovdqu32 zmm17, zmmword ptr [{target.arg64(5)} + 1 * 64]") |
| output.append(f"vpxord zmm1, zmm1, zmm17") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 1 * 64], zmm1") |
| output.append(f"vmovdqu32 zmm18, zmmword ptr [{target.arg64(5)} + 2 * 64]") |
| output.append(f"vpxord zmm2, zmm2, zmm18") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 2 * 64], zmm2") |
| output.append(f"vmovdqu32 zmm19, zmmword ptr [{target.arg64(5)} + 3 * 64]") |
| output.append(f"vpxord zmm3, zmm3, zmm19") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 3 * 64], zmm3") |
| output.append(f"vmovdqu32 zmm20, zmmword ptr [{target.arg64(5)} + 4 * 64]") |
| output.append(f"vpxord zmm4, zmm4, zmm20") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 4 * 64], zmm4") |
| output.append(f"vmovdqu32 zmm21, zmmword ptr [{target.arg64(5)} + 5 * 64]") |
| output.append(f"vpxord zmm5, zmm5, zmm21") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 5 * 64], zmm5") |
| output.append(f"vmovdqu32 zmm22, zmmword ptr [{target.arg64(5)} + 6 * 64]") |
| output.append(f"vpxord zmm6, zmm6, zmm22") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 6 * 64], zmm6") |
| output.append(f"vmovdqu32 zmm23, zmmword ptr [{target.arg64(5)} + 7 * 64]") |
| output.append(f"vpxord zmm7, zmm7, zmm23") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 7 * 64], zmm7") |
| output.append(f"vmovdqu32 zmm24, zmmword ptr [{target.arg64(5)} + 8 * 64]") |
| output.append(f"vpxord zmm8, zmm8, zmm24") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 8 * 64], zmm8") |
| output.append(f"vmovdqu32 zmm25, zmmword ptr [{target.arg64(5)} + 9 * 64]") |
| output.append(f"vpxord zmm9, zmm9, zmm25") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 9 * 64], zmm9") |
| output.append(f"vmovdqu32 zmm26, zmmword ptr [{target.arg64(5)} + 10 * 64]") |
| output.append(f"vpxord zmm10, zmm10, zmm26") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 10 * 64], zmm10") |
| output.append(f"vmovdqu32 zmm27, zmmword ptr [{target.arg64(5)} + 11 * 64]") |
| output.append(f"vpxord zmm11, zmm11, zmm27") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 11 * 64], zmm11") |
| output.append(f"vmovdqu32 zmm28, zmmword ptr [{target.arg64(5)} + 12 * 64]") |
| output.append(f"vpxord zmm12, zmm12, zmm28") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 12 * 64], zmm12") |
| output.append(f"vmovdqu32 zmm29, zmmword ptr [{target.arg64(5)} + 13 * 64]") |
| output.append(f"vpxord zmm13, zmm13, zmm29") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 13 * 64], zmm13") |
| output.append(f"vmovdqu32 zmm30, zmmword ptr [{target.arg64(5)} + 14 * 64]") |
| output.append(f"vpxord zmm14, zmm14, zmm30") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 14 * 64], zmm14") |
| output.append(f"vmovdqu32 zmm31, zmmword ptr [{target.arg64(5)} + 15 * 64]") |
| output.append(f"vpxord zmm15, zmm15, zmm31") |
| output.append(f"vmovdqu32 zmmword ptr [{target.arg64(5)} + 15 * 64], zmm15") |
| else: |
| raise NotImplementedError |
| else: |
| raise NotImplementedError |
| |
| |
| def xof_fn(target, output, degree, xor): |
| variant = "xor" if xor else "stream" |
| finish_fn_2d = xof_xor_finish_2d if xor else xof_stream_finish_2d |
| finish_fn_3d = xof_xor_finish_3d if xor else xof_stream_finish_3d |
| label = f"blake3_{target.extension}_xof_{variant}_{degree}" |
| output.append(f".global {label}") |
| output.append(f"{label}:") |
| if target.extension == AVX512: |
| if degree in (1, 2, 4): |
| xof_setup_2d(target, output, degree) |
| output.append(f"call {kernel_2d_name(target, degree)}") |
| finish_fn_2d(target, output, degree) |
| elif degree in (8, 16): |
| xof_setup_3d(target, output, degree) |
| output.append(f"call {kernel_3d_name(target, degree)}") |
| finish_fn_3d(target, output, degree) |
| else: |
| raise NotImplementedError |
| elif target.extension == AVX2: |
| assert degree == 2 |
| xof_setup_2d(target, output, degree) |
| output.append(f"call {kernel_2d_name(target, degree)}") |
| finish_fn_2d(target, output, degree) |
| elif target.extension in (SSE41, SSE2): |
| assert degree == 1 |
| xof_setup_2d(target, output, degree) |
| output.append(f"call {kernel_2d_name(target, degree)}") |
| finish_fn_2d(target, output, degree) |
| else: |
| raise NotImplementedError |
| output.append(target.ret()) |
| |
| |
| def emit_prelude(target, output): |
| # output.append(".intel_syntax noprefix") |
| pass |
| |
| |
| def emit_sse2(target, output): |
| target = replace(target, extension=SSE2) |
| kernel_2d(target, output, 1) |
| compress(target, output) |
| xof_fn(target, output, 1, xor=False) |
| xof_fn(target, output, 1, xor=True) |
| output.append(".balign 16") |
| output.append("PBLENDW_0x33_MASK:") |
| output.append(".long 0xFFFFFFFF, 0x00000000, 0xFFFFFFFF, 0x00000000") |
| output.append("PBLENDW_0xCC_MASK:") |
| output.append(".long 0x00000000, 0xFFFFFFFF, 0x00000000, 0xFFFFFFFF") |
| output.append("PBLENDW_0x3F_MASK:") |
| output.append(".long 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0x00000000") |
| output.append("PBLENDW_0xC0_MASK:") |
| output.append(".long 0x00000000, 0x00000000, 0x00000000, 0xFFFFFFFF") |
| |
| |
| def emit_sse41(target, output): |
| target = replace(target, extension=SSE41) |
| kernel_2d(target, output, 1) |
| compress(target, output) |
| xof_fn(target, output, 1, xor=False) |
| xof_fn(target, output, 1, xor=True) |
| |
| |
| def emit_avx2(target, output): |
| target = replace(target, extension=AVX2) |
| kernel_2d(target, output, 2) |
| xof_fn(target, output, 2, xor=False) |
| xof_fn(target, output, 2, xor=True) |
| |
| |
| def emit_avx512(target, output): |
| target = replace(target, extension=AVX512) |
| |
| # degree 1 |
| kernel_2d(target, output, 1) |
| compress(target, output) |
| xof_fn(target, output, 1, xor=False) |
| xof_fn(target, output, 1, xor=True) |
| |
| # degree 2 |
| kernel_2d(target, output, 2) |
| xof_fn(target, output, 2, xor=False) |
| xof_fn(target, output, 2, xor=True) |
| |
| # degree 4 |
| kernel_2d(target, output, 4) |
| xof_fn(target, output, 4, xor=False) |
| xof_fn(target, output, 4, xor=True) |
| |
| # degree 16 |
| kernel_3d(target, output, 16) |
| xof_fn(target, output, 16, xor=False) |
| xof_fn(target, output, 16, xor=True) |
| |
| |
| def emit_footer(target, output): |
| output.append(".balign 16") |
| output.append("BLAKE3_IV:") |
| output.append("BLAKE3_IV_0:") |
| output.append(".long 0x6A09E667") |
| output.append("BLAKE3_IV_1:") |
| output.append(".long 0xBB67AE85") |
| output.append("BLAKE3_IV_2:") |
| output.append(".long 0x3C6EF372") |
| output.append("BLAKE3_IV_3:") |
| output.append(".long 0xA54FF53A") |
| |
| output.append(".balign 16") |
| output.append("ROT16:") |
| output.append(".byte 2, 3, 0, 1, 6, 7, 4, 5, 10, 11, 8, 9, 14, 15, 12, 13") |
| output.append("ROT8:") |
| output.append(".byte 1, 2, 3, 0, 5, 6, 7, 4, 9, 10, 11, 8, 13, 14, 15, 12") |
| |
| output.append(".balign 64") |
| output.append("INCREMENT_2D:") |
| output.append(".quad 0, 0, 1, 0, 2, 0, 3, 0") |
| output.append("INCREMENT_3D:") |
| output.append(".quad 0, 1, 2, 3, 4, 5, 6, 7") |
| |
| output.append(".balign 64") |
| output.append("EVEN_INDEXES:") |
| output.append(".long 0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30") |
| output.append("ODD_INDEXES:") |
| output.append(".long 1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29, 31") |
| |
| |
| def format(output): |
| print("# This file is generated by asm.py. Don't edit this file directly.") |
| for item in output: |
| if ":" in item or item[0] == ".": |
| print(item) |
| else: |
| print(" " * 8 + item) |
| |
| |
| def main(): |
| target = Target(os=LINUX, arch=X86_64, extension=None) |
| output = [] |
| |
| emit_prelude(target, output) |
| emit_sse2(target, output) |
| emit_sse41(target, output) |
| emit_avx2(target, output) |
| emit_avx512(target, output) |
| emit_footer(target, output) |
| |
| format(output) |
| |
| |
| if __name__ == "__main__": |
| main() |
