Improve SHA2 precompile (#1480)

* Improve SHA2 precompile

* Review

* Add removed global label
This commit is contained in:
Robin Salen 2024-01-22 08:02:38 -05:00 committed by GitHub
parent 8e1969db2f
commit 319fc6a2f2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 81 additions and 45 deletions

View File

@ -4,6 +4,9 @@
// stack: num_blocks
%mul_const(320)
%add_const(2)
PUSH @SEGMENT_KERNEL_GENERAL
GET_CONTEXT
%build_address
%endmacro
global sha2_compression:
@ -24,9 +27,7 @@ global sha2_compression:
// stack: i=0, message_schedule_addr, a[0]..h[0], retdest
SWAP1
// stack: message_schedule_addr, i=0, a[0]..h[0], retdest
PUSH 0
// stack: 0, message_schedule_addr, i=0, a[0]..h[0], retdest
%mload_current_general
%mload_current_general_no_offset
// stack: num_blocks, message_schedule_addr, i=0, a[0]..h[0], retdest
DUP1
// stack: num_blocks, num_blocks, message_schedule_addr, i=0, a[0]..h[0], retdest
@ -53,7 +54,7 @@ compression_loop:
// stack: 4*i, message_schedule_addr, a[i], b[i], c[i], d[i], e[i], f[i], g[i], h[i], num_blocks, scratch_space_addr, message_schedule_addr, i, a[0]..h[0], retdest
ADD
// stack: message_schedule_addr + 4*i, a[i], b[i], c[i], d[i], e[i], f[i], g[i], h[i], num_blocks, scratch_space_addr, message_schedule_addr, i, a[0]..h[0], retdest
%mload_current_general_u32
%mload_u32
// stack: W[i], a[i], b[i], c[i], d[i], e[i], f[i], g[i], h[i], num_blocks, scratch_space_addr, message_schedule_addr, i, a[0]..h[0], retdest
PUSH sha2_constants_k
// stack: sha2_constants_k, W[i], a[i], b[i], c[i], d[i], e[i], f[i], g[i], h[i], num_blocks, scratch_space_addr, message_schedule_addr, i, a[0]..h[0], retdest

View File

@ -1,20 +1,22 @@
global sha2:
// stack: virt, num_bytes, retdest
SWAP1
// stack: num_bytes, virt, retdest
DUP2
// stack: virt, num_bytes, virt, retdest
%mstore_current_general
// stack: virt, retdest
PUSH @SEGMENT_KERNEL_GENERAL
GET_CONTEXT
%build_address
// stack: addr, num_bytes, retdest
DUP1 SWAP2
// stack: num_bytes, addr, addr, retdest
MSTORE_GENERAL
// stack: addr, retdest
// Precodition: input is in memory, starting at virt of kernel general segment, of the form
// Precondition: input is in memory, starting at addr of kernel general segment, of the form
// num_bytes, x[0], x[1], ..., x[num_bytes - 1]
// Postcodition: output is in memory, starting at 0, of the form
// num_blocks, block0[0], ..., block0[63], block1[0], ..., blocklast[63]
global sha2_pad:
// stack: virt, retdest
%mload_current_general
// stack: addr, retdest
MLOAD_GENERAL
// stack: num_bytes, retdest
// STEP 1: append 1
// insert 128 (= 1 << 7) at x[num_bytes+1]
@ -50,8 +52,7 @@ global sha2_pad:
DUP1
// stack: num_blocks, num_blocks, retdest
// STEP 5: write num_blocks to x[0]
PUSH 0
%mstore_current_general
%mstore_current_general_no_offset
// stack: num_blocks, retdest
%message_schedule_addr_from_num_blocks
%jump(sha2_gen_all_message_schedules)

View File

@ -3,9 +3,12 @@
// stack: num_blocks
%mul_const(64)
%add_const(2)
PUSH @SEGMENT_KERNEL_GENERAL
GET_CONTEXT
%build_address
%endmacro
// Precodition: stack contains address of one message block, followed by output address
// Precondition: stack contains address of one message block, followed by output address
// Postcondition: 256 bytes starting at given output address contain the 64 32-bit chunks
// of message schedule (in four-byte increments)
gen_message_schedule_from_block:
@ -16,18 +19,17 @@ gen_message_schedule_from_block:
// stack: block_addr + 32, block_addr, output_addr, retdest
SWAP1
// stack: block_addr, block_addr + 32, output_addr, retdest
%mload_current_general_u256
%mload_u256
// stack: block[0], block_addr + 32, output_addr, retdest
SWAP1
// stack: block_addr + 32, block[0], output_addr, retdest
%mload_current_general_u256
%mload_u256
// stack: block[1], block[0], output_addr, retdest
SWAP2
// stack: output_addr, block[0], block[1], retdest
%add_const(28)
PUSH 8
// stack: counter=8, output_addr + 28, block[0], block[1], retdest
%jump(gen_message_schedule_from_block_0_loop)
gen_message_schedule_from_block_0_loop:
// Split the first half (256 bits) of the block into the first eight (32-bit) chunks of the message sdchedule.
// stack: counter, output_addr, block[0], block[1], retdest
@ -43,7 +45,7 @@ gen_message_schedule_from_block_0_loop:
// stack: block[0] % (1 << 32), block[0] >> 32, output_addr, counter, block[1], retdest
DUP3
// stack: output_addr, block[0] % (1 << 32), block[0] >> 32, output_addr, counter, block[1], retdest
%mstore_current_general_u32
%mstore_u32
// stack: block[0] >> 32, output_addr, counter, block[1], retdest
SWAP1
// stack: output_addr, block[0] >> 32, counter, block[1], retdest
@ -81,7 +83,7 @@ gen_message_schedule_from_block_1_loop:
// stack: block[1] % (1 << 32), block[1] >> 32, output_addr, counter, block[0], retdest
DUP3
// stack: output_addr, block[1] % (1 << 32), block[1] >> 32, output_addr, counter, block[0], retdest
%mstore_current_general_u32
%mstore_u32
// stack: block[1] >> 32, output_addr, counter, block[0], retdest
SWAP1
// stack: output_addr, block[1] >> 32, counter, block[0], retdest
@ -111,39 +113,43 @@ gen_message_schedule_remaining_loop:
// stack: counter, output_addr, block[0], block[1], retdest
SWAP1
// stack: output_addr, counter, block[0], block[1], retdest
DUP1
// stack: output_addr, output_addr, counter, block[0], block[1], retdest
%sub_const(8)
PUSH 8
DUP2
// stack: output_addr, 2*4, output_addr, counter, block[0], block[1], retdest
SUB
// stack: output_addr - 2*4, output_addr, counter, block[0], block[1], retdest
%mload_current_general_u32
%mload_u32
// stack: x[output_addr - 2*4], output_addr, counter, block[0], block[1], retdest
%sha2_sigma_1
// stack: sigma_1(x[output_addr - 2*4]), output_addr, counter, block[0], block[1], retdest
SWAP1
// stack: output_addr, sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
DUP1
// stack: output_addr, output_addr, sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
%sub_const(28)
PUSH 28
DUP2
// stack: output_addr, 7*4, output_addr, sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
SUB
// stack: output_addr - 7*4, output_addr, sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
%mload_current_general_u32
%mload_u32
// stack: x[output_addr - 7*4], output_addr, sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
SWAP1
// stack: output_addr, x[output_addr - 7*4], sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
DUP1
// stack: output_addr, output_addr, x[output_addr - 7*4], sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
%sub_const(60)
PUSH 60
DUP2
// stack: output_addr, 15*4, output_addr, x[output_addr - 7*4], sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
SUB
// stack: output_addr - 15*4, output_addr, x[output_addr - 7*4], sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
%mload_current_general_u32
%mload_u32
// stack: x[output_addr - 15*4], output_addr, x[output_addr - 7*4], sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
%sha2_sigma_0
// stack: sigma_0(x[output_addr - 15*4]), output_addr, x[output_addr - 7*4], sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
SWAP1
// stack: output_addr, sigma_0(x[output_addr - 15*4]), x[output_addr - 7*4], sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
DUP1
// stack: output_addr, output_addr, sigma_0(x[output_addr - 15*4]), x[output_addr - 7*4], sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
%sub_const(64)
PUSH 64
DUP2
// stack: output_addr, 16*4, output_addr, sigma_0(x[output_addr - 15*4]), x[output_addr - 7*4], sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
SUB
// stack: output_addr - 16*4, output_addr, sigma_0(x[output_addr - 15*4]), x[output_addr - 7*4], sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
%mload_current_general_u32
%mload_u32
// stack: x[output_addr - 16*4], output_addr, sigma_0(x[output_addr - 15*4]), x[output_addr - 7*4], sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
SWAP1
// stack: output_addr, x[output_addr - 16*4], sigma_0(x[output_addr - 15*4]), x[output_addr - 7*4], sigma_1(x[output_addr - 2*4]), counter, block[0], block[1], retdest
@ -155,7 +161,7 @@ gen_message_schedule_remaining_loop:
// stack: sigma_1(x[output_addr - 2*4]) + x[output_addr - 16*4] + sigma_0(x[output_addr - 15*4]) + x[output_addr - 7*4], output_addr, counter, block[0], block[1], retdest
DUP2
// stack: output_addr, sigma_1(x[output_addr - 2*4]) + x[output_addr - 16*4] + sigma_0(x[output_addr - 15*4]) + x[output_addr - 7*4], output_addr, counter, block[0], block[1], retdest
%mstore_current_general_u32
%mstore_u32
// stack: output_addr, counter, block[0], block[1], retdest
%add_const(4)
// stack: output_addr + 4, counter, block[0], block[1], retdest
@ -178,12 +184,14 @@ global sha2_gen_all_message_schedules:
// stack: output_addr, retdest
DUP1
// stack: output_addr, output_addr, retdest
PUSH 0
// stack: 0, output_addr, output_addr, retdest
%mload_current_general
%mload_current_general_no_offset
// stack: num_blocks, output_addr, output_addr, retdest
PUSH 1
// stack: cur_addr = 1, counter = num_blocks, output_addr, output_addr, retdest
// stack: cur_offset = 1, counter = num_blocks, output_addr, output_addr, retdest
PUSH @SEGMENT_KERNEL_GENERAL
GET_CONTEXT
%build_address
// stack: cur_addr, counter, output_addr, output_addr, retdest
gen_all_message_schedules_loop:
// stack: cur_addr, counter, cur_output_addr, output_addr, retdest
PUSH gen_all_message_schedules_loop_end

View File

@ -1,5 +1,8 @@
%macro sha2_write_length
// stack: last_addr, length
// stack: last_addr_offset, length
PUSH @SEGMENT_KERNEL_GENERAL
GET_CONTEXT
%build_address
SWAP1
// stack: length, last_addr
DUP1
@ -8,7 +11,7 @@
// stack: length % (1 << 8), length, last_addr
DUP3
// stack: last_addr, length % (1 << 8), length, last_addr
%mstore_current_general
%swap_mstore
%rep 7
// For i = 0 to 6
@ -26,7 +29,7 @@
// stack: (length >> (8 * (i + 1))) % (1 << 8), length >> (8 * (i + 1)), last_addr - i - 2
DUP3
// stack: last_addr - i - 2, (length >> (8 * (i + 1))) % (1 << 8), length >> (8 * (i + 1)), last_addr - i - 2
%mstore_current_general
%swap_mstore
%endrep
%pop2

View File

@ -124,6 +124,16 @@
// stack: value
%endmacro
// Load a single value from the kernel general memory, in the current context (not the kernel's context).
%macro mload_current_general_no_offset
// stack:
PUSH @SEGMENT_KERNEL_GENERAL
GET_CONTEXT
%build_address_no_offset
MLOAD_GENERAL
// stack: value
%endmacro
// Load a big-endian u32 from kernel general memory in the current context.
%macro mload_current_general_u32
// stack: offset
@ -185,6 +195,19 @@
// stack: (empty)
%endmacro
// Store a single value to kernel general memory in the current context.
%macro mstore_current_general_no_offset
// stack: value
PUSH @SEGMENT_KERNEL_GENERAL
// stack: segment, value
GET_CONTEXT
// stack: context, segment, value
%build_address_no_offset
SWAP1
MSTORE_GENERAL
// stack: (empty)
%endmacro
%macro mstore_current_general(offset)
// stack: value
PUSH $offset