Improving integration of quantum_info module and qibojit

[renatomello]

Right now, even though the quantum_info module is compatible to all backends, it is not making use of any of the jit properties from numba. Therefore, numpy and numba performances are the same across the board. It would be good to profile the module and see where it is worth to invest time to code custom operations in qibojit.

Tasks

Preview Give feedback

1. Parallelization of CliffordOperations through numba and cupy qibojit#161
   enhancement +1
   
2. Customizing Clifford operations for qibojit engines #1150
   enhancement quantum_info module +2
   
3. Optimizations for the qibo.quantum_info.basis.pauli_basis and vectorization function #1459
   documentation enhancement quantum_info module +3
   

[renatomello]

@AlejandroSopena @BrunoLiegiBastonLiegi

This is list of the functions that have priority in being optimized in qibojit:

Basis:

• comp_basis_to_pauli
• pauli_to_comp_basis

Random ensembles:

• random_statevector
• random_density_matrix
• random_unitary
• random_quantum_channel
• random_clifford

Superoperator transformations:

• vectorization
• unvectorization
• _reshuffliing
• to_pauli_liouville
• choi_to_kraus
• kraus_to_choi
• kraus_to_stinespring
• stinespring_to_kraus

[BrunoLiegiBastonLiegi]

@renatomello should we create a separate module in quantum_info, something like _utils.py and create the relative modules in qibojit for numba and cupy for this?
Then I thought we could delegate to the backend the loading of this utils through a method. For example in NumbaBackend:

@staticmethod
def _load_quantum_info_utils():
    from qibojit.quantum_info import _utils_cpu
    
    return _utils_cpu

and then the functions inside the quantum_info module will use the provided backend to load the utils and call them.

[renatomello]

@renatomello should we create a separate module in quantum_info, something like _utils.py and create the relative modules in qibojit for numba and cupy for this? Then I thought we could delegate to the backend the loading of this utils through a method. For example in NumbaBackend:

@staticmethod
def _load_quantum_info_utils():
    from qibojit.quantum_info import _utils_cpu
    
    return _utils_cpu

and then the functions inside the quantum_info module will use the provided backend to load the utils and call them.

Which routines would be there?

[BrunoLiegiBastonLiegi]

Which routines would be there?

All the ones we want to improve with numba and cupy

[renatomello]

Which routines would be there?

All the ones we want to improve with numba and cupy

Then they'd have to be subroutines of the functions we have, because the functions we want to improve are all user-facing and part of specific submodules

[BrunoLiegiBastonLiegi]

@AlejandroSopena @BrunoLiegiBastonLiegi

This is list of the functions that have priority in being optimized in qibojit:

Basis:

* [ ]  `comp_basis_to_pauli`

* [ ]  `pauli_to_comp_basis`

Random ensembles:

* [ ]  `random_statevector`

* [ ]  `random_density_matrix`

* [ ]  `random_unitary`

* [ ]  `random_quantum_channel`

* [ ]  `random_clifford`

Superoperator transformations:

* [ ]  `vectorization`

* [ ]  `unvectorization`

* [ ]  `_reshuffliing`

* [ ]  `to_pauli_liouville`

* [ ]  `choi_to_kraus`

* [ ]  `kraus_to_choi`

* [ ]  `kraus_to_stinespring`

* [ ]  `stinespring_to_kraus`

#1459 is trying to address the first two boxes (basis section)

[BrunoLiegiBastonLiegi]

@renatomello I've done a quick test of random_statevector, by jitting with numba:

@njit("c16[:](i8)", parallel=True, cache=True)
def numba_random_statevec(dims):

    state = np.random.standard_normal(dims)
    state = state + 1.0j * np.random.standard_normal(dims)
    state = state / np.linalg.norm(state)

    return state

and explicitely using cupy:

def cupy_random_statevec(dims):

    state = cp.random.standard_normal(dims)
    state = state + 1.0j * cp.random.standard_normal(dims)
    state = state / cp.linalg.norm(state)

    return state

I was able to test only up to 25 qubits due to my local memory limitations and I got these results with 10 repetitions each:

######## random statevector ########
-------- 5 qubits -------
numpy:  0.0005540169986488763
numba:  0.0014047330005269032
cupy:  0.16188096400219365
-------- 15 qubits -------
numpy:  0.011496602000988787
numba:  0.09011737599939806
cupy:  0.16356382600133657
-------- 25 qubits -------
numpy:  10.051634777999425
numba:  2.770201025999995
cupy:  0.15904992800096807

The only small issue is that we lose direct access to the seed, which should be then reset globally with np.random.seed/cp.random.seed.

[renatomello]

@BrunoLiegiBastonLiegi these custom functions can only accept one input?

[BrunoLiegiBastonLiegi]

For cupy not necessarily but for numba, yes, I had to drop it because np.random.default_rng was not working.

[renatomello]

I guess we could drop seed just for numba and then mention in the docstrings that this is a numba issue, not a qibo issue.

EDIT: Maybe this helps https://stackoverflow.com/questions/55806542/how-to-compile-a-numba-jited-function-with-variable-input-type ?

[BrunoLiegiBastonLiegi]

Yeah i mean you could always have a wrapper:

@njit("c16[:](i8)", parallel=True, cache=True)
def _random_statevector(dims):

    state = np.random.standard_normal(dims)
    state = state + 1.0j * np.random.standard_normal(dims)
    state = state / np.linalg.norm(state)

    return state

def random_statevector(dims, seed):
    np.random.seed(seed)
    return _random_staevector(dims)

I am more worried on how to integrate this and similar changes to the backend, to be honest...

[renatomello]

I am more worried on how to integrate this and similar changes to the backend, to be honest...

If you don't know, then I really do not know.

[BrunoLiegiBastonLiegi]

I mean most likely we will need to make random_statevector part of the backend, and similarly for many other functions in the module, but, firstly, we were discussing lately about reducing the size of the backend in terms of what it does, and secondly, this would probably require tearing apart a significant fraction of the quantum_info module.
At least, for what concerns numba integration,. For cupy it is relatively easier as long as you don't write custom cuda kernels. The problem with cupy is that its np attribute is numpy rather than cupy which is fine when you call:

backend.np.function(cp.array)

because the cupy implementation is triggered anyway, but not, as for example in this case, when you do sampling or creation:

backend.np.random.random_normal()

[alecandido]

In a sense, you may want to do for the quantum_info module what we were planning to do for the most part of the current backends: split the state manipulation from the circuit execution itself.

However, since the other plan is not advancing at all right now, if you want to make the quantum_info module integrated with the various backends, then just expand them, as much as needed.
It is true that this will make the current situation even worse - but I'd argue that it won't be considerably worse, and instead mostly similar to what is already happening.

---

Alternative proposal explanation/reminder

What we were already discussing (even in the scope of qibo-core) was to split the state manipulation.
The sense was that what mostly differs in the various backends is the execution of the circuit itself, which is what defines the backend (consider also Qibolab as an example). Then, you're left with an output, which may have many alternative features

• full state (density matrix or state vector or partial observable) vs shots (for some observable)
• the place where is stored: CPU vs GPU
• derivative availability

However, multiple backends may share the same outcome features. E.g. NumPy, TensorFlow, and Qibolab may all produce an observable made of shots on the CPU memory. And they can be manipulated in the very same way. Same for the full state vector on GPU, which is independent of the library who generated it (it may be recasted as a buffer, if needed).
So, the proposal was to keep any manipulation function separate, to avoid duplication. Just defining the standard structures which can be manipulated.

However, this plan requires some considerable refactoring effort, and it's not going to happen immediately.

[BrunoLiegiBastonLiegi]

Apart from the separation between the execution and manipulation backends, which is probably the best solution but takes some time, I was thinking about providing the backends with a backend.quantum_info attribute where to store all the custom methods that we'd like to customize. This way the integration of the quantum_info module and the backends becomes easier, while still maintaining a reasonable level of isolation. Namely, it's quite easy to roll back and have everything under control since it's all localized in one place. What do you guys think?

[alecandido]

@BrunoLiegiBastonLiegi to me, there is no problem.

As argued above, I know is very much against isolation, but given the current situation, I'd be willing to accept even to proliferate directly in the scope of the backend itself, adding methods there.

Instead, if you instead believe that you can reuse this inner scopes (the backend.quantum_info attribute) in multiple backends (and especially if they are not already related by inheritance in the way which is most convenient to you), I'd encourage you to proceed with your plan.

Otherwise, it's just the same :)