Open ZzEeKkAa opened 1 year ago
Hi @ZzEeKkAa,
Thank you for tring the feature.
We can reproduce this on linux
Yes, you right. There is some unimplemented interface. jaxlib.xla_extension.XlaRuntimeError: UNIMPLEMENTED: PJRT C API does not support AcquireExternalReference
[stream=0x55e8dc9039e0,impl=0x55e8dd530ff0] StreamPool returning ok stream
2023-07-07 14:45:37.326709: I itex/core/compiler/xla/service/gpu/gpu_executable.cc:477] GpuExecutable::ExecuteAsyncOnStreamImpl(jit_iota) time: 129 ms (cumulative: 129 ms, max: 129 ms, #called: 1)
2023-07-07 14:45:37.326719: I itex/core/compiler/xla/pjrt/pjrt_stream_executor_client.cc:1989] Replica 0 partition 0 completed; ok=1
2023-07-07 14:45:37.326884: I itex/core/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2268] Replicated execution complete.
Traceback (most recent call last):
File "/ws1/xigui/itex-source-code/dlpack.py", line 2, in <module>
jnp.arange(10, dtype=jnp.float32).__dlpack__()
File "/home/xiguiwang/ws1/conda/itex-build/lib/python3.9/site-packages/jax/_src/array.py", line 343, in __dlpack__
return to_dlpack(self)
File "/home/xiguiwang/ws1/conda/itex-build/lib/python3.9/site-packages/jax/_src/dlpack.py", line 51, in to_dlpack
return xla_client._xla.buffer_to_dlpack_managed_tensor(
jaxlib.xla_extension.XlaRuntimeError: UNIMPLEMENTED: PJRT C API does not support AcquireExternalReference
2023-07-07 14:45:37.338990: I itex/core/compiler/xla/pjrt/pjrt_stream_executor_client.cc:1208] PjRtStreamExecutorBuffer::Delete
2023-07-07 14:45:37.339017: I itex/core/compiler/xla/stream_executor/stream_executor_pimpl.cc:217] Called StreamExecutor::Deallocate(mem=0xffff81ac00000000) mem->size()=40
2023-07-07 14:45:37.350224: I itex/core/compiler/xla/stream_executor/stream_executor_pimpl.cc:367] Called StreamExecutor::SynchronizeAllActivity()
@yiqianglee @jzhoulon
Do you have idea this "jaxlib.xla_extension.XlaRuntimeError: UNIMPLEMENTED: PJRT C API does not support AcquireExternalReference"?
Is this a problem of "PJRT C API" Interface design of NOT cover/support or a problem of ITEX NOT implementaton?
Xigui
Yes, this is a problem of "PJRT C API" Interface design and we need support from OpenXLA side. https://github.com/openxla/xla/blob/main/xla/pjrt/pjrt_c_api_client.h#L281-L285
@ZzEeKkAa Hope this helps for you.
This works on recent Intel extension for Tensorflow 2.15 + oneAPI 2024.1.
This is the software version I used to verify this problem. intel-extension-for-tensorflow 2.15.0.0 pypi_0 pypi intel-extension-for-tensorflow-lib 2.15.0.0.2 pypi_0 pypi jax 0.4.26 pypi_0 pypi jaxlib 0.4.26 pypi_0 pypi
Here is the document to set up your intel-extension-for-tensorflow: https://intel.github.io/intel-extension-for-tensorflow/v2.15.0.0/docs/install/install_for_xpu.html#intel-xpu-software-installation
System: WSL2 Ubuntu 22.04, on top of Windows 11 CPU: 1270P GPU: integrated (
[ext_oneapi_level_zero:gpu:0] Intel(R) Level-Zero, Intel(R) Graphics [0x46a6] 1.3 [1.3.26032]
) Tensorflow: 2.12.0 Jax: 0.4.4Results into error:
Jax itself works on level zero GPU, so environment is not broken. I guess it is lack of implementation of
PJRT C API does not support AcquireExternalReference
. It blocks from users workflows that require both jax related operations and, for example, numba_dpex related operations without memory copying.