blas/mkl issue on Google colab

[twiecki]

Here's an example google colab that installs aesara 2.7.7 and then tries to import: https://colab.research.google.com/drive/1bPq3mcIlcw6mJQE8Gf8bcEdVhJdJSUpB#scrollTo=36fiA6TJjLJC

Error:

---------------------------------------------------------------------------
NoSectionError                            Traceback (most recent call last)
[/usr/local/lib/python3.7/dist-packages/aesara/configparser.py](https://localhost:8080/#) in fetch_val_for_key(self, key, delete_key)
    236             try:
--> 237                 return self._aesara_cfg.get(section, option)
    238             except InterpolationError:

12 frames
NoSectionError: No section: 'blas'

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
KeyError: 'blas__ldflags'

During handling of the above exception, another exception occurred:

ModuleNotFoundError                       Traceback (most recent call last)
ModuleNotFoundError: No module named 'mkl'

During handling of the above exception, another exception occurred:

RuntimeError                              Traceback (most recent call last)
[/usr/local/lib/python3.7/dist-packages/aesara/link/c/cmodule.py](https://localhost:8080/#) in check_mkl_openmp()
   2691 you set this flag and don't set the appropriate environment or make
   2692 sure you have the right version you *will* get wrong results.
-> 2693 """
   2694         )
   2695 

RuntimeError: 
Could not import 'mkl'.  If you are using conda, update the numpy
packages to the latest build otherwise, set MKL_THREADING_LAYER=GNU in
your environment for MKL 2018.

If you have MKL 2017 install and are not in a conda environment you
can set the Aesara flag blas__check_openmp to False.  Be warned that if
you set this flag and don't set the appropriate environment or make
sure you have the right version you *will* get wrong results.

Setting %env MKL_THREADING_LAYER=GNU allows importing, but it's not clear why aesara can't figure that out itself, or what the problem is exactly.

Output of check_blas.py:

        Some results that you can compare against. They were 10 executions
        of gemm in float64 with matrices of shape 2000x2000 (M=N=K=2000).
        All memory layout was in C order.

        CPU tested: Xeon E5345(2.33Ghz, 8M L2 cache, 1333Mhz FSB),
                    Xeon E5430(2.66Ghz, 12M L2 cache, 1333Mhz FSB),
                    Xeon E5450(3Ghz, 12M L2 cache, 1333Mhz FSB),
                    Xeon X5560(2.8Ghz, 12M L2 cache, hyper-threads?)
                    Core 2 E8500, Core i7 930(2.8Ghz, hyper-threads enabled),
                    Core i7 950(3.07GHz, hyper-threads enabled)
                    Xeon X5550(2.67GHz, 8M l2 cache?, hyper-threads enabled)


        Libraries tested:
            * numpy with ATLAS from distribution (FC9) package (1 thread)
            * manually compiled numpy and ATLAS with 2 threads
            * goto 1.26 with 1, 2, 4 and 8 threads
            * goto2 1.13 compiled with multiple threads enabled

                          Xeon   Xeon   Xeon  Core2 i7    i7     Xeon   Xeon
        lib/nb threads    E5345  E5430  E5450 E8500 930   950    X5560  X5550

        numpy 1.3.0 blas                                                775.92s
        numpy_FC9_atlas/1 39.2s  35.0s  30.7s 29.6s 21.5s 19.60s
        goto/1            18.7s  16.1s  14.2s 13.7s 16.1s 14.67s
        numpy_MAN_atlas/2 12.0s  11.6s  10.2s  9.2s  9.0s
        goto/2             9.5s   8.1s   7.1s  7.3s  8.1s  7.4s
        goto/4             4.9s   4.4s   3.7s  -     4.1s  3.8s
        goto/8             2.7s   2.4s   2.0s  -     4.1s  3.8s
        openblas/1                                        14.04s
        openblas/2                                         7.16s
        openblas/4                                         3.71s
        openblas/8                                         3.70s
        mkl 11.0.083/1            7.97s
        mkl 10.2.2.025/1                                         13.7s
        mkl 10.2.2.025/2                                          7.6s
        mkl 10.2.2.025/4                                          4.0s
        mkl 10.2.2.025/8                                          2.0s
        goto2 1.13/1                                                     14.37s
        goto2 1.13/2                                                      7.26s
        goto2 1.13/4                                                      3.70s
        goto2 1.13/8                                                      1.94s
        goto2 1.13/16                                                     3.16s

        Test time in float32. There were 10 executions of gemm in
        float32 with matrices of shape 5000x5000 (M=N=K=5000)
        All memory layout was in C order.


        cuda version      8.0    7.5    7.0
        gpu
        M40               0.45s  0.47s
        k80               0.92s  0.96s
        K6000/NOECC       0.71s         0.69s
        P6000/NOECC       0.25s

        Titan X (Pascal)  0.28s
        GTX Titan X       0.45s  0.45s  0.47s
        GTX Titan Black   0.66s  0.64s  0.64s
        GTX 1080          0.35s
        GTX 980 Ti               0.41s
        GTX 970                  0.66s
        GTX 680                         1.57s
        GTX 750 Ti               2.01s  2.01s
        GTX 750                  2.46s  2.37s
        GTX 660                  2.32s  2.32s
        GTX 580                  2.42s
        GTX 480                  2.87s
        TX1                             7.6s (float32 storage and computation)
        GT 610                          33.5s
        
Some Aesara flags:
    blas__ldflags= -L/usr/local/lib -lmkl_rt -lpthread -lm -lm
    compiledir= /root/.aesara/compiledir_Linux-5.4.188+-x86_64-with-Ubuntu-18.04-bionic-x86_64-3.7.13-64
    floatX= float64
    device= cpu
Some OS information:
    sys.platform= linux
    sys.version= 3.7.13 (default, Apr 24 2022, 01:04:09) 
[GCC 7.5.0]
    sys.prefix= /usr
Some environment variables:
    MKL_NUM_THREADS= None
    OMP_NUM_THREADS= None
    GOTO_NUM_THREADS= None

Numpy config: (used when the Aesara flag "blas__ldflags" is empty)
blas_mkl_info:
  NOT AVAILABLE
blis_info:
  NOT AVAILABLE
openblas_info:
    libraries = ['openblas', 'openblas']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None)]
    runtime_library_dirs = ['/usr/local/lib']
blas_opt_info:
    libraries = ['openblas', 'openblas']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None)]
    runtime_library_dirs = ['/usr/local/lib']
lapack_mkl_info:
  NOT AVAILABLE
openblas_lapack_info:
    libraries = ['openblas', 'openblas']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None)]
    runtime_library_dirs = ['/usr/local/lib']
lapack_opt_info:
    libraries = ['openblas', 'openblas']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', 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
Numpy dot module: numpy
Numpy location: /usr/local/lib/python3.7/dist-packages/numpy/__init__.py
Numpy version: 1.21.6

We executed 10 calls to gemm with a and b matrices of shapes (5000, 5000) and (5000, 5000).

Total execution time: 74.73s on CPU (with direct Aesara binding to blas).

Try to run this script a few times. Experience shows that the first time is not as fast as following calls. The difference is not big, but consistent.

CC @canyon289 @junpenglao @ColCarroll

[canyon289]

Thanks for summarizing,
Are there any identified fixes that don't involve modifying an environment variable?

[brandonwillard]

It looks like the mkl-service package isn't installed, because import mkl fails. Installing mkl-service leads to the following error:

Python 3.7.13
---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
[<ipython-input-6-fed9734531b7>](https://localhost:8080/#) in <module>()
      1 get_ipython().system('python --version')
----> 2 import mkl

[/usr/local/lib/python3.7/dist-packages/mkl/__init__.py](https://localhost:8080/#) in <module>()
     47 
     48 with RTLD_for_MKL():
---> 49     from . import _mklinit
     50 
     51 del RTLD_for_MKL

ImportError: libmkl_rt.so.2: cannot open shared object file: No such file or directory

---------------------------------------------------------------------------
NOTE: If your import is failing due to a missing package, you can
manually install dependencies using either !pip or !apt.

To view examples of installing some common dependencies, click the
"Open Examples" button below.
---------------------------------------------------------------------------

In other words, the environment doesn't appear to have the requisite MKL shared libraries.

The reason it's even assuming there's an MKL installation is due to the following:

import numpy.distutils.system_info

numpy.distutils.system_info.get_info("blas_opt")
# {'define_macros': [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)],
#  'include_dirs': ['/usr/local/include'],
#  'libraries': ['mkl_rt', 'pthread'],
#  'library_dirs': ['/usr/local/lib']}

This numpy.distutils approach was removed in #1050. The new approach would yield the following:

import numpy as np

np.__config__.get_info("blas_opt")
# {'define_macros': [('HAVE_CBLAS', None)],
#  'language': 'c',
#  'libraries': ['openblas', 'openblas'],
#  'library_dirs': ['/usr/local/lib'],
#  'runtime_library_dirs': ['/usr/local/lib']}

I'm not sure which configuration is correct, but at least the one produced by the new approach appears to avoid MKL entirely.

[twiecki]

Thanks for taking a look. So I guess 2.7.7 does not contain #1050 yet? Can we do a new release and try with that?

As for the env: it seems like no MKL is present, not sure why numpy/scipy thinks it does, but defaulting to openblas here is probably the safest option.

[twiecki]

I just tried with aesara 2.7.8 and it works. So I think we just need a new pymc version with updated deps and we'll be good to go.

[canyon289]

Thats perfect, thanks so much @brandonwillard and @twiecki for figuring this out, I'll add the new release of pymc to the colab container