EnzymeAD / EnzymeAD/Enzyme-JAX

Export issue with missing BLAS symbol

Open
#135 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
MLIR
Stars
131
Forks
53
Avg merge
1d 10h
Merged PRs (30d)
193

Description

Ran into the following error when trying to use the enzyme_jax export functionality:

Note, we found this bug when running the pip installed version of jax so v0.0.8

```
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1727977776.038651 4451776 service.cc:148] XLA service 0x393930100 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
I0000 00:00:1727977776.038733 4451776 service.cc:156] StreamExecutor device (0): Host, Default Version
2024-10-03 10:49:41.235014: E external/xla/xla/service/cpu/simple_orc_jit.cc:453] Unable to resolve runtime symbol: `_blas_dtrsm'. Hint: if the symbol a custom call target, make sure you've registered it with the JIT using XLA_CPU_REGISTER_CUSTOM_CALL_TARGET.
2024-10-03 10:49:41.235054: E external/xla/xla/service/cpu/simple_orc_jit.cc:453] Unable to resolve runtime symbol: `_xla_python_cpu_callback'. Hint: if the symbol a custom call target, make sure you've registered it with the JIT using XLA_CPU_REGISTER_CUSTOM_CALL_TARGET.
JIT session error: Symbols not found: [ _xla_python_cpu_callback, _blas_dtrsm ]
Traceback (most recent call last):
File "ExaCMech/pyecmech/jax_ecmech/jax_evptn_wrap.py", line 314, in
stress_vec_pressure_n1, history_update, internal_energy_n1, temp_k, sdd, junk = evptn_wc.solve(
^^^^^^^^^^^^^^^
File "ExaCMech/pyecmech/jax_ecmech/jax_evptn_wrap.py", line 239, in solve
export("./get_response_jit.ir", self.get_response_jit,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "anaconda3/lib/python3.11/site-packages/enzyme_ad/jax/primitives.py", line 1562, in export
enzyme_call.compile_to_llvm(
ValueError: NOT_FOUND: Symbol _main.20658 not found.

```

A reproducer for this can be found when trying to run this python script here: https://github.com/LLNL/ExaCMech/blob/57829e133966c39b7b731fbbfc4efe448fa36400/pyecmech/jax_ecmech/jax_evptn_wrap.py

The README here https://github.com/LLNL/ExaCMech/blob/jax_export/pyecmech/jax_ecmech/README.md also covers the required pip installable packages required to run things. Although, one can ignore the note about needing to install the general library that this python library exists in.

I can provide more about my local install but I imagine these version numbers should be helpful info:
```
conda 24.5.0
enzyme_ad 0.0.8
equinox 0.11.4
jax 0.4.33
jaxlib 0.4.33
jaxtyping 0.2.29
numpy 1.26.4
opt-einsum 3.3.0
optimistix 0.0.7
```

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.