Implement SWAP fusion

[TysonRayJones]

Note

For all development, please work out of the devel branch

Context

It may be worth a quick, prior read about QuEST's software architecture.

SWAP

The SWAP gate is ubiquitous in quantum computing and effects a change upon the quantum state as if the targeted qubits were physically swapped or relabelled.

Let $i_{[t]}$ notate the $t$-th bit of an unsigned integer $i$, and let $i_{\neg {a,b}}$ notate the unsigned integer resulting from flipping two bits of $i$; those at indices $a$ and $b$. As outlined here (Sec IV C), the SWAP gate targeting qubits $t_1$ and $t_2$ modifies the $i$-th basis state $\ket{i}$ as

$$ SWAP_{t_1,t_2} \ket{i} = \begin{cases} \ket{i} \text{ when } i_{[t_1]} = i_{[t_2]}, \text{ else } \ket{i_{\neg {t_1,t_2}}} \end{cases}$$

When $SWAP_{t_1,t_2}$ operates upon a statevector $\ket{\psi} = \sum\limits_i \alpha_i \ket{i}$, only amplitudes $\alpha_i$ which satisfy $i_{[t_1]} \ne i_{[t_2]}$ are changed; they are swapped with a pair amplitude at index $i\neg{t_1,t_2}$.

$$ \alpha_i \rightarrow \alpha_{i\neg{t_1,t_2}} $$

QuEST's code to simulate a SWAP gate in non-distributed settings resembles:

QuEST/quest/src/cpu/cpu_subroutines.cpp

Lines 258 to 282 in 53f8f3a

	template <int NumCtrls>
	void cpu_statevec_anyCtrlSwap_subA(Qureg qureg, vector<int> ctrls, vector<int> ctrlStates, int targ1, int targ2) {
	assert_numCtrlsMatchesNumCtrlStatesAndTemplateParam(ctrls.size(), ctrlStates.size(), NumCtrls);
	// each control qubit halves the number of iterations, each of which modifies 2 amplitudes, and skips 2
	qindex numIts = qureg.numAmpsPerNode / powerOf2(2 + ctrls.size());
	auto sortedQubits = util_getSorted(ctrls, {targ2, targ1});
	auto qubitStateMask = util_getBitMask(ctrls, ctrlStates, {targ2, targ1}, {0, 1});
	// use template param to compile-time unroll loop in insertBits()
	SET_VAR_AT_COMPILE_TIME(int, numCtrlBits, NumCtrls, ctrls.size());
	int numQubitBits = numCtrlBits + 2;
	#pragma omp parallel for if(qureg.isMultithreaded)
	for (qindex n=0; n<numIts; n++) {
	// i01 = nth local index where ctrls are active, targ2=0 and targ1=1
	qindex i01 = insertBitsWithMaskedValues(n, sortedQubits.data(), numQubitBits, qubitStateMask);
	qindex i10 = flipTwoBits(i01, targ2, targ1);
	std::swap(qureg.cpuAmps[i01], qureg.cpuAmps[i10]);
	}
	}

Read how to simulate this operation in distributed settings here (Sec IV C).

multi-SWAP

Many quantum algorithms feature circuits containing a contiguous sequence of non-overlapping SWAP gates.

Some unrelated QuEST routines even use a sequence of SWAPs to reorder qubits, in order to change the communication pattern during distributed simulation.

This is called "cache blocking" and is used by distributed simulators like cuQuantum, Qiskit, mpiQulacs, often as the primary means of distribution. In QuEST, this technique is necessary when effecting arbitrary matrices targeting two or more qubits, like via applyCompMatr2. See Sec. IV D (pg 21) here for an explanation.

It is prudent to simulate a multi-SWAP (i.e. a sequence of contiguous SWAP gates) as fast as possible, and to reduce its communication costs when simulated in distributed settings.

Problem

Currently, when QuEST needs to perform a multi-SWAP, it simply performs each individual SWAP in-turn:

QuEST/quest/src/core/localiser.cpp

Lines 906 to 932 in 53f8f3a

	void anyCtrlMultiSwapBetweenPrefixAndSuffix(Qureg qureg, vector<int> ctrls, vector<int> ctrlStates, vector<int> targsA, vector<int> targsB) {
	// this is an internal function called by the below routines which require
	// performing a sequence of SWAPs to reorder qubits, or move them into suffix.
	// the SWAPs act on unique qubit pairs and so commute.
	/// @todo
	/// - the sequence of pair-wise full-swaps should be more efficient as a
	/// "single" sequence of smaller messages sending amps directly to their
	/// final destination node. This could use a new "multiSwap" function.
	/// - if the user has compiled cuQuantum, and Qureg is GPU-accelerated, the
	/// multiSwap function should use custatevecSwapIndexBits() if local,
	/// or custatevecDistIndexBitSwapSchedulerSetIndexBitSwaps() if distributed,
	/// although the latter requires substantially more work like setting up
	/// a communicator which may be inelegant alongside our own distribution scheme.
	// perform necessary swaps to move all targets into suffix, each of which invokes communication
	for (size_t i=0; i<targsA.size(); i++) {
	if (targsA[i] == targsB[i])
	continue;
	int suffixTarg = std::min(targsA[i], targsB[i]);
	int prefixTarg = std::max(targsA[i], targsB[i]);
	anyCtrlSwapBetweenPrefixAndSuffix(qureg, ctrls, ctrlStates, suffixTarg, prefixTarg);
	}
	}

This is inefficient! Each successive SWAP is not modifying any amplitudes, but merely moving them. It should be much faster to move amplitude each straight to their final locations, performing all swaps "at once". This is called "SWAP fusion", first reported here, and requires more complicated communication logic than the single-SWAP case. SWAP fusion is implemented by e.g. cuQuantum and mpiQulacs.

Solution

Optimising multi-SWAP via implementing SWAP fusion would greatly benefit QuEST's distributed simulation of many-target matrix gates and Kraus maps (which use the aforementioned reordering trick). Implementing SWAP fusion will involve replacing the one-by-one swaps in localiser.cpp's anyCtrlMultiSwapBetweenPrefixAndSuffix with a new, custom, distributed routine in comm_routines.cpp. It may also require locally-parallelised (OpenMP and CUDA) routines to pack communication buffers, implemented in cpu_subroutines.cpp and gpu_subroutines.cpp.

An external contribution need not attempt implementing the GPU logic; the OpenMP logic alone is fine.

Implementing this will also require benchmarking to confirm the fused SWAPs are faster, including when each distributed node is not parallelised (with multithreading or GPU-acceleration), to ensure the communication benefits are not outweighed by additional local processing/branching (which is of course very unlikely).

[TysonRayJones]

Noting PR #643 drafted an incomplete solution

[Amitav-Krishna]

Hey, I'm interested in taking a crack at this, is this exclusively for people participating in UnitaryHack?

[TysonRayJones]

Hi there,
You're very welcome to to contribute this feature - it's open!

[Amitav-Krishna]

Hey, getting started on this, will have a draft PR done tomorrow. I'm planning on creating a virtual-physical qubit mapping, so that when we do swaps instead of transferring gigabytes of data over the network we can swap the mappings. Before, SWAP(t_1, t_2) would physically move them around in memory, my change would just require changing the values to which the mapping, and then from there we could implement the fusion based on some heuristics. This would be especially helpful for when the states are on different nodes. Does this sound good? N

[TysonRayJones]

Hmm I'm apprehensive about adding that layer of indirection above the state, though I understand the idea and its benefits quite well*. Right now, I prefer QuEST's simplicity; the state is exactly "where" you expect, making custom integrations trivial. It's straightforward for software layers above QuEST (which don't offer direct amplitude access to its users) to implement wire ordering for a free speedup, while low-level integrations see no indirections, and the QuEST API/doc has no need to distinguish the concepts of "wires" and "qubits". Furthermore, a "wire ordering" indirection opens the door to a series of further related optimisations that significantly complicate the code base. Such an indirection might be too substantial a refactor for an external contirbution.

I could be overruled on this though, if e.g. @otbrown has different opinions on the balance to strike between simplicity and performance. For now, it's certainly possible to implement a fused multi-SWAP operation which doesn't necessitate the indirection (i.e. the fused SWAP operation is performed immediately, moving amplitudes).

*I presently work on cuStateVec which has a "wire ordering" indirection explained here

[Amitav-Krishna]

Ah, okay, sorry I guess I didn't really understand the layer QuEST was working at well enough, I will go with the more direct method of just directly fusing their swap operations.

[TysonRayJones]

Oh no need at all to be sorry - it's a good idea, and something potentially worth debating about!

[Amitav-Krishna]

Hey, going to implement it now. Will open a draft PR tonight, I'm going to just implement the more trivial fusions (SWAP(x,y), SWAP(x, y), SWAP(x, y) -> SWAP(x, y)) so that I can get a feel for the codebase. Tomorrow I'll implement the multi-qubit swap fusions (SWAP(x, y), SWAP(y, z), SWAP(x, y) -> SWAP(x, z)). For now I won't do any of the wire indirection either,

[TysonRayJones]

Just noting that neither of the scenarios you mention are a "multi-SWAP". In the illustrations above, the targets of the sequence of the SWAP gates are all unique. That's critical; we are swapping one group of qubits with a disjoint group. It's not a matter of cancelling or simplifying a sequence of SWAPs.

[Amitav-Krishna]

Ah, so this is more about parallelizing the communication? Right now the problem is that we execute the communication sequentially, and the ideal PR would change this so that it would all be in parallel, packing all of the buffers and sending all of the MPI requests at once, overlapping them. Am I interpreting this correctly now? Sorry for misunderstanding earlier

[TysonRayJones]

Closer, but not quite. Right now, each quantum SWAP operation is performed in-turn. Some amplitudes moved between nodes by one SWAP operation will be wastefully moved again by a subsequent operation. It is possible to avoid this by simulating all SWAP operations as a single "fused" operation, sending amplitudes directly to their ultimate nodes