add support for avx512 and allow users to choose the Stim architecture

[MoiseRousseau]

Hi!

I started implementing the AVX512F instructions for code vectorization, which is more recent and is available on newer AMD Zen4 and on some Intel processors. This can give up to 4% improved performance over AVX2 version on my very limited testing.

Also, I added in the Python binding the possibility to the user to specify the architecture he/she wants to use with the STIM_ARCHITECTURE environment variable. The default behavior is still SSE2 because of #432.

Below some basic benchmark (1% improvement in this case):

(env_dev) moise@moise-HP-EliteBook-840:~/Program/Stim/examples$ export STIM_ARCH=avx512 && time python surface_code.py 
real	0m32.448s
user	0m26.335s
sys	0m6.786s

(env_dev) moise@moise-HP-EliteBook-840:~/Program/Stim/examples$ export STIM_ARCH=avx2 && time python surface_code.py 
real	0m32.787s
user	0m26.680s
sys	0m6.782s

(env_dev) moise@moise-HP-EliteBook-840:~/Program/Stim/examples$ export STIM_ARCH=sse2 && time python surface_code.py 
real	0m35.860s
user	0m29.884s
sys	0m6.656s

with the code:

(env_dev) moise@moise-HP-EliteBook-840:~/Program/Stim/examples$ cat surface_code.py 
import stim 

c = stim.Circuit.generated(
"surface_code:rotated_memory_z",
rounds=20,
distance=40,
after_clifford_depolarization=0.001,
after_reset_flip_probability=0.001,
before_measure_flip_probability=0.001,
before_round_data_depolarization=0.001)

sampler = c.compile_sampler()
samples = sampler.sample(300000)

TODO: correct the function inplace_transpose_block_pass_64_and_128_and_256 line 234 in [src/stim/mem/simd_word_512_avx512.h](https://github.com/quantumlib/Stim/compare/main...MoiseRousseau:Stim:Moise/arch_avx512?expand=1#diff-75e1a1cb88702d8bf8b9f964584cababa9c07f6868db23d08a5b1015b3410ae3). However, it seems this function is not called in the code..

Let me know what you think,
Moise

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[Strilanc]

Cool change! It looks like it was relatively straightforward to implement.

The main thing missing here is there are no unit tests. Look for the definition of TEST_EACH_WORD_SIZE_W in the code and add the 512 case, so that the 512 simd word is unit tested, verifying the methods work correctly. It's possible you will run into issues where some tests presupposed that the max bit width was 256, and need to be tweaked.

Unfortunately, we can't yet update the python extension to use avx because of the mysterious collab crash that it causes. I'm scared that enabling avx will silently crash some unknown subset of users, with no way for me to notice it is happening, and I don't yet understand what is actually causing the crash.

(I'd also ask for you to add benchmarks, but the work to templatize the simulators so you can ask for a 512 or 256 simd simulation is not yet done so these benchmarks are essentially impossible to define for the moment.)

[Strilanc]

Unfortunately, due to the collab bug, this still can't be uncommented.

[MoiseRousseau]

Hi @Strilanc !

I will have a look to the unit tests for 512 bits word.

Regarding the crash, I cannot reproduce it on colab with stim version 1.10.0 (published October 25th, so before the fix on November 19th 2022). On a new and fresh notebook, I enter the following:

!pip install stim==1.10.0
import stim
stim.Circuit().shortest_graphlike_error()

which returns:

Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/
Collecting stim==1.10.0
  Downloading stim-1.10.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.7 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 4.7/4.7 MB 993.6 kB/s eta 0:00:00
Requirement already satisfied: numpy in /usr/local/lib/python3.9/dist-packages (from stim==1.10.0) (1.22.4)
Installing collected packages: stim
Successfully installed stim-1.10.0

---------------------------------------------------------------------------

ValueError                                Traceback (most recent call last)

[<ipython-input-1-5e5a8bf45752>](https://localhost:8080/#) in <cell line: 3>()
      1 get_ipython().system('pip install stim==1.10.0')
      2 import stim
----> 3 stim.Circuit().shortest_graphlike_error()

ValueError: Failed to find any graphlike logical errors.

Do you still have the crash ?

Moise

[Strilanc]

Even if it doesn't reproduce on collab anymore... what if it still happens on other people's machines? That's the terrifying part.

[MoiseRousseau]

Yeah agree, but I was asking to try have a simple look on what could happen. Because what if this crash also occurs with the AVX512 implementation ?

[Strilanc]

@MoiseRousseau Even if the crash occurs with avx512, we can still add the implementation and turn it on later.

[Strilanc]

I'm going to close this as stale. Still open to something like this being added.