"Illegal instruction" crash in AVX512 code from XLA when AVX512 is not supported by a Xen hypervisor · Issue #2906 · jax-ml/jax (original) (raw)

Script:

fav25@dev-gpu-0:~$ python3 Python 3.6.9 (default, Apr 18 2020, 01:56:04) [GCC 8.4.0] on linux Type "help", "copyright", "credits" or "license" for more information.

import jax key = jax.random.PRNGKey(0) /auto/homes/fav25/jax/jax/lib/xla_bridge.py:116: UserWarning: No GPU/TPU found, falling back to CPU. warnings.warn('No GPU/TPU found, falling back to CPU.') jax.random.normal(key, shape=(3,))

Output

Illegal instruction (core dumped)

I have followed the installation instructions for CPU using pip and also the installation instructions building from source (cloned the latest version git clone https://github.com/google/jax):

python build/build.py 
pip install -e build  # install jaxlib (includes XLA)
pip install -e .      # install jax (pure Python)

Linux version:

$ lsb_release -a
Distributor ID: Ubuntu
Description:    Ubuntu 18.04.4 LTS
Release:        18.04
Codename:       bionic

I could not find the file _pywrap_xla.so searched inside the build directory and in the jax folder since the package was installed from there it wont be in site-packages.

 $ find . | grep "_pywrap_xla.so"
fav25@dev-gpu-0:~/jax$                                 

Most similar binary I found was:

 $ ldd xla_extension.so
        linux-vdso.so.1 (0x00007ffd981c2000)
        libstdc++.so.6 => /usr/lib/x86_64-linux-gnu/libstdc++.so.6 (0x00007fe050339000)
        libm.so.6 => /lib/x86_64-linux-gnu/libm.so.6 (0x00007fe04ff9b000)
        libdl.so.2 => /lib/x86_64-linux-gnu/libdl.so.2 (0x00007fe04fd97000)
        libpthread.so.0 => /lib/x86_64-linux-gnu/libpthread.so.0 (0x00007fe04fb78000)
        libgcc_s.so.1 => /lib/x86_64-linux-gnu/libgcc_s.so.1 (0x00007fe04f960000)
        libc.so.6 => /lib/x86_64-linux-gnu/libc.so.6 (0x00007fe04f56f000)
        /lib64/ld-linux-x86-64.so.2 (0x00007fe053f3a000)

The outcomes are the same when installed using pip.

lscpu output:

$ lscpu
Architecture:        x86_64
CPU op-mode(s):      32-bit, 64-bit
Byte Order:          Little Endian
CPU(s):              8
On-line CPU(s) list: 0-7
Thread(s) per core:  1
Core(s) per socket:  1
Socket(s):           8
NUMA node(s):        1
Vendor ID:           GenuineIntel
CPU family:          6
Model:               85
Model name:          Intel(R) Xeon(R) Gold 6142 CPU @ 2.60GHz
Stepping:            4
CPU MHz:             2594.136
BogoMIPS:            5187.88
Hypervisor vendor:   Xen
Virtualisation type: full
L1d cache:           32K
L1i cache:           32K
L2 cache:            1024K
L3 cache:            22528K
NUMA node0 CPU(s):   0-7