Implement multithreaded stochastic swap in rust

[mtreinish]

Summary

This commit is a rewrite of the core swap trials functionality in the
StochasticSwap transpiler pass. Previously this core routine was written
using Cython (see #1789) which had great performance, but that
implementation was single threaded. The core of the stochastic swap
algorithm by it's nature is well suited to be executed in parallel, it
attempts a number of random trials and then picks the best result
from all the trials and uses that for that layer. These trials can
easily be run in parallel as there is no data dependency between the
trials (there are shared inputs but read-only). As the algorithm
generally scales exponentially the speed up from running the trials in
parallel can offset this somewhat and improve the general scaling of the
pass to a point. Running the pass in parallel was previously tried in #4781
using Python multiprocessing but the overhead of launching an additional
process and serializing the input arrays for each trial was significantly larger
than the speed gains. To run the algorithm efficiently in parallel
multithreading is needed to leverage shared memory on shared inputs.

This commit rewrites the cython routine using rust. This was done for
two reasons. The first is that rust's safety guarantees make dealing
with and writing parallel code much easier and safer. It's also
multiplatform because the rust language supports native threading
primitives in language. The second is while writing parallel cython
code using open-mp there are limitations with it, mainly on windows. In
practice it was also difficult to write and maintain parallel cython
code as it has very strict requirements on python and C code
interactions. It was much faster and easier to port it to rust and the
performance for each iteration (outside of parallelism) is the same (and in
some cases marginally faster) in rust. The implementation here reuses
the data structures that the previous cython implementation introduced
(mainly flattening all the terra objects into 1d or 2d numpy arrays for
efficient access from C).

The speedups from this PR can be significant, calling transpile() on a
400 qubit (with a depth of 10) QV model circuit targeting a 409 heavy
hex coupling map goes from ~200 seconds with the single threaded cython
to ~60 seconds with this PR locally on a 32 core system, When transpiling
a 1000 qubit (also with a depth of 10) QV model circuit targeting a 1081
qubit heavy hex coupling map goes from taking ~3700 seconds to ~720
seconds.

The tradeoff with this PR is for local qiskit-terra development a rust
compiler needs to be installed. This is made trivial using rustup
(https://rustup.rs/), but it is an additional burden and one that we
might not want to make. If so we can look at turning this PR into a
separate repository/package that qiskit-terra can depend on. The
tradeoff here is that we'll be adding friction to the api boundary
between the pass and the core swap trials interface. But, it does ease
the dependency on development for qiskit-terra.

Details and comments

Fixes #1743

TODO:

• Fix test failures
• Fix hanging test/test timeout
• Fix tox configuration and docs build
• Add Docs
• Update contributing docs to include requiring a rust compiler
• Add release notes
• More thorough benchmarking and potentially determining if there is a break even point for parallel vs serial (and if there is use that as a default switch)
• Update linux wheel publish jobs to install rust compiler

[mtreinish]

Just for comparison I reran the 1k qubit QV transpile with logging enabled per this script:

import logging
import time

import numpy as np
import qiskit
from qiskit import *
from qiskit.quantum_info import random_unitary
from qiskit.transpiler import CouplingMap

logging.basicConfig(level="INFO")


def build_qv_model_circuit(width, depth, seed=None):
    """
    The model circuits consist of layers of Haar random
    elements of SU(4) applied between corresponding pairs
    of qubits in a random bipartition.
    """
    np.random.seed(seed)
    circuit = QuantumCircuit(width)
    # For each layer
    for _ in range(depth):
        # Generate uniformly random permutation Pj of [0...n-1]
        perm = np.random.permutation(width)
        # For each pair p in Pj, generate Haar random SU(4)
        for k in range(int(np.floor(width/2))):
            U = random_unitary(4)
            pair = int(perm[2*k]), int(perm[2*k+1])
            circuit.append(U, [pair[0], pair[1]])
    return circuit

qv = build_qv_model_circuit(1000, 10)
basis_gates = ["id", "u1", "u2", "u3", "cx"]
cmap = CouplingMap.from_heavy_hex(21)
start_time = time.time()
transpile(qv, coupling_map=cmap, basis_gates=basis_gates, seed_transpiler=42)
stop_time = time.time()
print(stop_time - start_time)

to get the isolated stochastic swap times as part of the larger transpile() call. With this PR applied it logged 553533.51545 ms and without this PR (using the serial cython implementation) it logged 3585305.39703 ms

[ewinston]

I profiled random circuits of depth 10 and variable width on a grid lattice. For the single threaded cython I got,

While for the multithreaded rust binary on a cpu with 16 cores I got a ~3.8x improvement in speed.

I noticed that while the rust version was running, although all the cores were in use they did not always seem to be at 100% which might partially account for the ratio.

[jakelishman]

I personally am fine with adding Rust components to Terra, but I have a much larger tolerance for increases in complexity of development tooling than most others, I think. It probably wouldn't hurt for us to coalesce our compiled components into either Cython or Rust, but not both, in due course.

I'd like to hear others' opinions on adding Rust (unless you all talked about this in a meeting while I was out), since I think it affects the larger Terra dev team in quite a big way.

The changes largely look fine to me - it looks like a pretty mechanical conversion of the Cython code to Rust (you've still got the variables ii, jj and kk which are a dead giveaway that Paul originally wrote it!). I left a few comments about typos, and a couple that are basically just beginner Rust questions.

[jakelishman]

This is minor, but something I was thinking about. Given floating-point maths, fold can't assume that the addition is associative, so it can't do fast-math-like optimisations like turning the fold into a threaded reduction. Given that this seems to be the worker routine of the loop (by eye - I did exactly zero profiling), is it worth us looking at that for large sizes?

The addition technically still isn't associative, but this shouldn't be a situation where we care about small differences - anything with a coupled link large enough to really screw with the calculation results will never be the best solution anyway, and there can't be catastrophic cancellation because the values should all be positive.

[mtreinish]

It is something I've thought about and why I switched from the for loop in the original cython to the iterator based approach here. My concern with doing a multithreaded reduce here was we're already in a parallel context when this is run so I was worried we just be introducing more overhead.

Basically if you do:

Suggested change

	#[inline]
	fn compute_cost(
	dist: &ArrayView2<f64>,
	layout: &NLayout,
	gates: &[usize],
	num_gates: usize,
	) -> f64 {
	(0..num_gates)
	.map(|kk| {
	let ii = layout.logic_to_phys[gates[2 * kk]];
	let jj = layout.logic_to_phys[gates[2 * kk + 1]];
	dist[[ii, jj]]
	})
	.fold(0.0, |a, b| a + b)
	}
	#[inline]
	fn compute_cost(
	dist: &ArrayView2<f64>,
	layout: &NLayout,
	gates: &[usize],
	num_gates: usize,
	) -> f64 {
	(0..num_gates)
	.into_par_iter()
	.map(|kk| {
	let ii = layout.logic_to_phys[gates[2 * kk]];
	let jj = layout.logic_to_phys[gates[2 * kk + 1]];
	dist[[ii, jj]]
	})
	.reduce(|| 0.0, |a, b| a + b)
	}

it becomes a parallel reduction. My other thought here was if there was maybe some simd operation we could use here too

As for the fp error here, yeah I agree. Honestly the typing was a bit fast and loose in the cython around this call. The distance matrix here is always an integer (although the array comes from retworkx distance_matrix() function) because it's just counting edges between nodes in the graph. So when this is called with cdist we don't really have to worry about it, but for the times this is called with scale that is actually a float.

[jakelishman]

To be honest, it probably needs a bunch of considerations for the fastest speed here, because there's a couple of levels of memory indirection before we get to the register variables we want.

Yeah, the threading inside process parallelisation may well not be worth it in the end. Maybe something to look at in the future, but probably not super high priority.

[jakelishman]

To be honest, it probably needs a bunch of considerations for the fastest speed here, because there's a couple of levels of memory indirection before we get to the register variables we want.

Yeah, the threading inside process parallelisation may well not be worth it in the end. Maybe something to look at in the future, but probably not super high priority.

[mtreinish]

Well we're already multithreaded here. That's where the speed advantage comes in this PR, the outer trials loop is executed in a rayon parallel iterator. This does raise a good point that we might still be in parallel processes too (although given the state of our multiprocessing support on windows, macOS>=py3.8, or linux==py3.9) and this multithreaded single pass might be more harmful for performance then. What I meant here though was nested multithreaded worker pools. I did try this locally just a little while ago to see what would happen and it basically has pegged all 64 threads on my workstation for about 1.5 hours now. So definitely not the approach we want to take. :) (we might be able to try to balance threads a bit like on my workstation I have 44 free threads, when running with 20 trials, which means we could probably get away with using 2 threads per trial here to do this sum in parallel, but my workstation is a pretty extreme case).

I was also looking at simdeez as a potential way to use simd here, but yeah I agree it will require more thought (I'll also have to brush up on my vectorized intrinsics to figure out how to use that lib) and is a future thing. We can leave optimizing this function as a future todo

[jakelishman]

Also, with us potentially adding a new compiled dependency, how does Rust handle choosing the set of CPU instructions it's going to target? I'm assuming it does something sensible, but just wanted to check that we're not accidentally going to start compiling in AVX-512 instructions if the build machine happens to have them available.

[mtreinish]

Also, with us potentially adding a new compiled dependency, how does Rust handle choosing the set of CPU instructions it's going to target? I'm assuming it does something sensible, but just wanted to check that we're not accidentally going to start compiling in AVX-512 instructions if the build machine happens to have them available.

By default it's fairly conservative. For x86_64 the only extra features over the core ISA it uses are fxsr, sse, and sse2 (which pretty much every 64bit x86 cpu supports). To get more than that you have to explicitly ask to tune for a specific CPU (either on the command line, via a env variable, or I think you can change the default in the Cargo.toml too). For example, you can do RUSTFLAGS='-Ctarget-cpu=native' pip install . and it will build for your native local CPU. We can obviously tune this to be something a bit more modern if we wanted, for example if we set the target cpu to x86_64-v2 it adds sse2, sse3, sse4.1, sse4.2, and popcount which basically limits us to everything newer than 2008. But for retworkx at least I opted to maximize compatibility in the precompiled binaries and just left it as is with the generic x86_64 target even if that potentially left some performance on the floor.

[mtreinish]

I noticed that while the rust version was running, although all the cores were in use they did not always seem to be at 100% which might partially account for the ratio.

The library I used for parallel iterators, rayon, uses a thread pool with work stealing to implement the parallel execution. When I was running on my 16 vCPU (pinned to 8 physical cores) windows vm and my old desktop/benchmarking machine (also with physical 8 cores and 16 threads) I saw a pretty clear pattern for optimization level 1 with 20 trials configured on each layer permutation that it would start full running on all the cores and then after those started to finish there were still 4 trials left which would run in 4 threads as the 16 initial executions started to finish. But as those last 4 ran most of the cores were idle. I didn't really look closer than that though, but if we're not getting 100% cpu utilization on each thread we can look at profiling this to see where things are waiting because I really wouldn't expect much idle time and if there is some that's pointing some further optimization we can do here.

[jlapeyre]

By default it's fairly conservative.

If this were a separate package, is it possible for the a user to optionally build from source while installing? The number of people willing to do it would be small enough that it is unlikely to be worth it unless there are already robust tools for installing from source.

[jlapeyre]

The tradeoff with this PR is for local qiskit-terra development a rust
compiler needs to be installed. This is made trivial using rustup
(https://rustup.rs/)

I mostly metoo @jakelishman 's first comment. I am inclined to vote yes. But, I use linux, so I find the burden is essentially zero. How many people develop Qiskit on windows? I often hear that installing and using tools on windows is not easy, even when it's superficially easy.

[mtreinish]

By default it's fairly conservative.

If this were a separate package, is it possible for the a user to optionally build from source while installing? The number of people willing to do it would be small enough that it is unlikely to be worth it unless there are already robust tools for installing from source.

Yeah, you can do this very easily either as a separate package or if it was part of qiskit-terra (there's functionally no difference for cpu tuning). Either way it's basically the same step you basically just need to set the env variable RUSTFLAGS=-Ctarget-cpu=native or RUSTFLAGS=-Ctarget-cpu=x86_64-v3 when you call pip install either on a local checkout or on the sdist we publish to pypi. That will be passed to the rust compiler internally by setuptools-rust when it compiles the rust extension. For example, on retworkx you can do this today with something like: RUSTFLAGS="-Ctarget-cpu=native" pip install --no-binary retworkx retworkx which will build it from the published sdist tuned for your local cpu.

[mtreinish]

The tradeoff with this PR is for local qiskit-terra development a rust
compiler needs to be installed. This is made trivial using rustup
(https://rustup.rs/)

I mostly metoo @jakelishman 's first comment. I am inclined to vote yes. But, I use linux, so I find the burden is essentially zero. How many people develop Qiskit on windows? I often hear that installing and using tools on windows is not easy, even when it's superficially easy.

I'm not sure how many people actively are developing on a windows as their daily driver. I do it from time to time, mostly to debug environment specific issues (I've lost count of the number of FP precision issues we've hit in windows), but it's not my daily/primary dev environment. I will say installing rust on Windows is about as easy as on linux. Instead of being a command line program your curl and run rustup is just a .exe which opens a shell and runs the same basic script (they also have a graphical .msi installer available, but I've never tried it). The only complexity is you do need a c compiler (well technically the linker) which on windows by default is msvc. But this was already a requirement to do terra dev with the cython stuff. So at least for me testing this PR on my VM yesterday I could just run the rustup .exe file I got from https://rustup.rs/ and everything just worked the same.

[jlapeyre]

will build it from the published sdist tuned for your local cpu.
That's really slick.

yesterday I could just run the rustup .exe

My comment was partly inspired by juliaup story (only partly, there are many other stories of installation woe on windows). juliaup is written in rust and is modeled closely on rustup. But, they still have problems. For example, winget install julia -s msstore can fail. But, it may be easier for developers to work around problems than average users. Maybe the juliaup authors have not yet been able to organize something as simple as a web page with a download link to rustup-init.exe.

Beyond installing rustup, there is getting rust to work. Searching shows things like this: https://stackoverflow.com/a/68835925 . I'm still not saying adding rust is the wrong move, just trying to predict what might happen... I suppose getting cython to work on windows can be a problem too, and eventually, that requirement could be dropped.

[jakelishman]

To be fair, that stackoverflow post is about linking Windows Rust against mingw's gcc rather than the first-class-supported Visual C++ tooling. That's a very niche use-case, and you'd have similar issues trying to force current Terra to build with a non-standard C compiler on any operating system.

"Getting Rust to work" shouldn't be an issue with pre-compiled wheels - it shouldn't be any different to linking against a compiled C library - Terra already relies on retworkx, which is all distributed Rust code.

[jlapeyre]

To be fair, that stackoverflow post is about linking Windows Rust against mingw's gcc

You are correct, they say explicitly that they are trying to do a non-standard installation. That didn't sink in the first time I read it. I could dig some more, but I'm not going to run into a show stopper. I guess the development requirements won't be an excessive burden. Starting with one module seems like a good idea. If it doesn't work out, it is not too much trouble to revert.