Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add dlpack support for XPU for JAX #9834

Closed
wants to merge 2 commits into from
Closed

Conversation

wozna
Copy link

@wozna wozna commented Feb 22, 2024

This PR adds support for dlpack for xpu (realated to #6151) for JAX with intel-extension-for-openxla:

  • adds platform name XpuName and platform_id XpuId
  • adds DLDeviceType for OneAPI (kDLOneAPI) and refators DeviceForDLDevice function
  • adds checking of environment variables (ONEAPI_DEVICE_SELECTOR, ZE_FLAT_DEVICE_HIERARCHY) required for proper dlpack communication between different frameworks. This is due to the fact that oneapi frameworks supporting dlpack use the sycl device list to identify devices. In JAX with intel_extesion_for_openxla, device identification does not use the entire sycl device list, but only a partial selection of "level_zero:gpu".
    A solution to use only part of this list is to use the ONEAPI_DEVICE_SELECTOR environment variable. Therefore, if the buffer is located on an XPU, it is necessary to check that the libraries communicating with each other use the same device list to set and read the device_id.
    If sycl is added to openxl it will be possible to remove this sanity check.

Generalize the env check function to all devices

Separate xpu checking func and refactor DeviceForDLDevice func
@github-actions github-actions bot added the kokoro:force-run Forces CI to rerun label Feb 22, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Forces CI to rerun label Feb 22, 2024
@github-actions github-actions bot added the kokoro:force-run Forces CI to rerun label Feb 23, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Forces CI to rerun label Feb 23, 2024
@Zantares
Copy link

Hi @wozna , please hold a while and check this PR #9042 first. You have implemented the DLPack feature based on old XPU name in Intel Extension, but we are trying to submit the common part to public OpenXLA and soon the device will be changed to GPU/SYCL. Some of the PR may be invalid soon. We can have a discussion of how to add DLPack feature in the new design, thanks!

@kamaljeeti
Copy link
Contributor

Hi @Zantares , any update on this PR? Thanks.

@Zantares
Copy link

Hi @Zantares , any update on this PR? Thanks.

I have discussed my concern with the PR author, and we figured out another way to solve her requirement. Please feel free to close this PR, thanks!

@kamaljeeti
Copy link
Contributor

Okay @Zantares, thanks for confirming. Closing this PR.

@kamaljeeti kamaljeeti closed this Mar 18, 2024
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.

6 participants