Skip to content

[OneAPI GPU] Enable hybrid solver kernels for the OneAPI backend - #40186

Open
kranipa wants to merge 1 commit into
jax-ml:mainfrom
Intel-tensorflow:kranipa/oneapi_hybrid_kernels
Open

[OneAPI GPU] Enable hybrid solver kernels for the OneAPI backend#40186
kranipa wants to merge 1 commit into
jax-ml:mainfrom
Intel-tensorflow:kranipa/oneapi_hybrid_kernels

Conversation

@kranipa

@kranipa kranipa commented Aug 25, 2026

Copy link
Copy Markdown
Contributor

This PR adds following changes:

  1. Wires the _hybrid extension into the OneAPI plugin.
  2. The check is added in SyclMemcpyAsync to synchronize the memcpy when usm buffer pointer type is usm::alloc::unknown. (D2H memory transfer.)
  3. jaxlib/gpu_solver.py - import and register the OneAPI _hybrid module under the ONEAPI platform.
  4. jax/_src/lax/linalg.py - register geqp3/eig lowerings for OneAPI.
  5. jaxlib/tools/build_gpu_kernels_wheel.py - package _hybrid.so into the OneAPI wheel.

@kranipa
kranipa marked this pull request as ready for review August 25, 2026 08:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant