Fix Loads and Stores for SHA256 benchmark

[aborg-dev]

This PR adds the remaining features necessary to run SHA256 benchmark:

• Misaligned loads and stores
• Loads and stores of size < 64 bit

There instructions are generated by rustc compiler for wasm32 target and are quite common.

The current implementation is not verifiable as it uses external (free) inputs. It turned out to be quite challenging to implement them in a verifiable fashion because they involve many bitwise operations which are quite expensive.

For now, there is value in having a non-verifiable implementation that makes sure we can run tests and benchmarks related to memory. I've linked a few TODOs to do a verifiable implementation, but first, we would need a proper design for that part. Most likely, we will need to modify the zkAsm processor to support these operations efficiently.

Implementation-wise, there are three steps:

• Conversion from address + offset to slot + offset
• Read/write the value at the correct offset
• Narrowing down the value to the desired type width

[nagisa]

Can you really rely on this? What happens if wasm goes like i32.load offset=0 (i32.const 7)? I think there's no choice but to check alignment of the final computed (byte) address at runtime and/or unconditionally load two slots.

[aborg-dev]

Fair point, I'm not sure whether I can rely on the fact that the address in the load/store register will be word-aligned. So we at least need an assert to check this.

Given that this is a non-trivial amount of work that we might need to redo anyway soon, I'll start with an assert in the code to check that the address is word-aligned. If it is not triggered by the Rust benchmarks that we use, that would be interesting evidence to investigate how exactly these addresses are generated. If it is triggered, we'll need to address it when we have the test/benchmark that exercises it.

[aborg-dev]

Well, this was very well spotted. I've added an assert and it triggered. I'll work on fixing it.

[nagisa]

FWIW for a minimal reproducer something like read_unaligned would do. In architectures that require aligned loads the compiler backend would responsibly split the load into aligned loads, but WebAssembly specifically does not impose that sort of requirement, so the backend doesn't need to worry about unaligned addresses at all.

[aborg-dev]

It took some effort, but I rewrote the code without the assumption that the address is word-aligned.
I ended up unconditionally reading and writing to two slots, though I think in the future it would not be that hard to optimize the code to skip the second write when possible using a jump.

[aborg-dev]

Converted this back to draft while I figure out the way to fix unaligned register addresses.

[nagisa]

This seems alright, although knowing zkasm, I think a quick improvement in steps in the future is going to involve JNZ on the result of % 8 and otherwise execute a single load only. Probably similar with stores as well.