How to parallelise contraction over slices?

[vincentmr]

Say I have some large circuit circ for which I want a local expectation value, but it could be anything really. I would like to use Cotengra to optimize the contraction cost and slicing. For example,

import cotengra as ctg

opt = ctg.ReusableHyperOptimizer(
    max_repeats=1024,
    methods=["greedy"],
    slicing_opts={"target_slices": 2**6},  # first do basic slicing
    max_time=60,
    parallel=True,
)
quimb_obs = qu.pauli("Z")
where = 0

rehs = circ.local_expectation_rehearse(quimb_obs, (where), optimize=opt, backend="jax")

Once I have rehs, how can I map slices to cores (assuming I have a multicore processor)?

[jcmgray]

cotengra has a few examples like this one: https://github.com/jcmgray/cotengra/blob/main/examples/ex_mpi_executor.py, the last lines of that script are most relevant. Just call:

tree = tn.contraction_tree(optimize=existing_path_or_optimizer)
arrays = tn.arrays

to get the 'cotengra' form of the contraction.

What I would note is that multi-threading or multiprocessing (basically shared memory parallelization) over slices is likely not beneficial. As long as the individual contractions are parallelized, then the unsliced contraction will require less memory and probably be at least as fast, (no slicing overhead, and as long as the contractions are big, the low level i.e. BLAS threading should be very efficient).

Slicing for parallelization in general only makes obvious sense for distributed memory, though maybe you have some specific reasons!

[vincentmr]

Thanks for the reference, this is what I was looking for.

As for the multi-threading, I'm getting started with tensor networks so I may have the wrong intuition here. I am playing with TNs derived from quantum circuits (with 1-qubit and 2-qubit gates), and hence I'm thinking that tensors and the corresponding BLAS operations should be fairly small and not scale that well on, say, an EPYC processor with 64 cores. While the total flops count will increase with slicing, the speed-up w.r.t. number of processes should be pretty close to perfect for expectation values where one only needs to reduce a scalar. I was going to make some test to confirm, but what do you think?

[jcmgray]

Yes I am not really a hardware expert, so hard to predict and certainly worth just trying instead.

I think for very thin contractions like that you quickly become memory speed ('write') bound, which maybe negates the extra parallelization.

Also if you use cotengra and minimize='combo' it will find contractions that minimize FLOPS + alpha * WRITE which gives much better practical performance (it works essentially like automated gate fusing) and increases the arithmetic intensity so that BLAS multithreading is still performant.

Likely there is some specific optimum combination of (num_blas_threads * num_slice_workers) == num_cpus for a given memory limit that is fastest for specific hardware.