EnzymeAD / EnzymeAD/Enzyme-JAX
Export issue with missing BLAS symbol
- 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.