qtn.Circuit.sample uses multiple CPU cores

[imbacalvin]

Hi Johnnie, while trying to sample several times from a very simple quantum circuit, I noticed that qtn.Circuit.sample uses all the CPU cores in my computer. I simply wonder is there a way to force qtn.Circuit.sample to be single-threaded? I am doing

import os
os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'
os.environ["NUMEXPR_NUM_THREADS"] = "1"

before importing numpy, so the multi-threading behavior probably doesn't come from numpy. I also tried setting

quimb.core._NUM_THREAD_WORKERS=1

which didn't seem to solve the problem.

Thanks!

[imbacalvin]

It appears that setting group_size=1 solves the problem, and the problem shows up if setting group_size>=8.

[jcmgray]

Hi @imbacalvin,

The only thing I can think of is that the contraction optimizer might be using multiple cores. The default one ("auto-hq" preset) should not however.

If you supply a minimal reproducing example I am happy to look into it at some point.

[imbacalvin]

Thanks for the quick reply!

I was able to reproduce the problem on my desktop using the example below. For this specific case, the problem goes away if I set group_size=13...

import quimb
import quimb.tensor as qtn

qc=qtn.circuit.CircuitMPS(60)
for i in range(0,60,2):
    qc.apply_gate_raw(quimb.gen.rand.rand_uni(4),[i,i+1])
for i in range(1,59,2):
    qc.apply_gate_raw(quimb.gen.rand.rand_uni(4),[i,i+1])
    
for x in qc.sample(1000,group_size=14):
    pass

Please let me know if you need further information!

[jcmgray]

I couldn't really pin down where the over-threading was taking place (wasn't very consistent), my guess is something to do with the auto-hq optimizer from cotengra but in principle that should be single threaded too. Using optimize="greedy" seems to fix things.

I've actually just added much more efficient MPS-specific sampling that side steps the issue here : #244.

But would still be good to work out exactly which bit of code is using the threads.

[imbacalvin]

Hi Johnnie, thanks for the quick update! The problem seems not to be tied with sampling specifically as the following example also invokes the multi-threading behavior.

import quimb
import quimb.tensor as qtn

qc=qtn.circuit.CircuitMPS(60)
for qq in range(10):
    for i in range(0,60,2):
        qc.apply_gate_raw(quimb.gen.rand.rand_uni(4),[i,i+1])
    for i in range(1,59,2):
        qc.apply_gate_raw(quimb.gen.rand.rand_uni(4),[i,i+1])

Multi-threading seems to come from apply_gate_raw. Is there a simple fix to make the above single threaded?

[jcmgray]

Hmm, this runs single threaded for me. What OS, numpy+BLAS are you using? Perhaps it is something not covered by the env vars above.

[imbacalvin]

Interesting... I'm running on Ubuntu 23.10 with numpy 1.24.4. numpy.__config__.show() returns the following

openblas64__info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None)]
    runtime_library_dirs = ['/usr/local/lib']
blas_ilp64_opt_info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None)]
    runtime_library_dirs = ['/usr/local/lib']
openblas64__lapack_info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None), ('HAVE_LAPACKE', None)]
    runtime_library_dirs = ['/usr/local/lib']
lapack_ilp64_opt_info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None), ('HAVE_LAPACKE', None)]
    runtime_library_dirs = ['/usr/local/lib']
Supported SIMD extensions in this NumPy install:
    baseline = SSE,SSE2,SSE3
    found = SSSE3,SSE41,POPCNT,SSE42,AVX,F16C,FMA3,AVX2
    not found = AVX512F,AVX512CD,AVX512_KNL,AVX512_KNM,AVX512_SKX,AVX512_CLX,AVX512_CNL,AVX512_ICL

Somehow for me, running the script below still uses all the CPU cores...

import os
os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["VECLIB_MAXIMUM_THREADS"] = "1"

import quimb
import quimb.tensor as qtn
quimb.core._NUM_THREAD_WORKERS=1

qc=qtn.circuit.CircuitMPS(60)
for qq in range(10):
    for i in range(0,60,2):
        qc.apply_gate_raw(quimb.gen.rand.rand_uni(4),[i,i+1])
    for i in range(1,59,2):
        qc.apply_gate_raw(quimb.gen.rand.rand_uni(4),[i,i+1])

[jcmgray]

Yes a bit mysterious. Is this system python or a conda/mamba env etc? You could try something like: https://github.com/joblib/threadpoolctl to see what thread pools are being launched.

[imbacalvin]

Thanks for your insights! The script above has become single-threaded for me as well, after not touching it for a few days and maybe after rebooting my computer? 🤷