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

[Draft][XLA:GPU] Add sycl runtime #9042

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

ShengYang1
Copy link
Contributor

@ShengYang1 ShengYang1 commented Feb 1, 2024

This PR aims to add sycl runtime support, we can run basic JAX GPU UTs with it. It includes:

  1. sycl runtime crosstool build
  2. sycl stream executor
  3. spirv-llvm-translator
  4. registration of spir_compiler/computation placer/transfer manager
  5. refine build file to remove cuda specific changes(especially those cuda kernels)
  6. fix CudaCubin to load spirv binary
  7. fix metadata

This is a draft PR to show what is needed to run JAX on sycl runtime.

@github-actions github-actions bot added the kokoro:force-run Forces CI to rerun label Feb 1, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Forces CI to rerun label Feb 1, 2024
@ShengYang1 ShengYang1 marked this pull request as draft February 1, 2024 06:54
@Zantares
Copy link

Zantares commented Feb 1, 2024

Hi @penpornk , this is the big SYCL runtime PR for preview and CI test, we will continue to split it into small PRs.

Next work:

We still need some JAX work to enable Intel public CI for this PR. Will share CI result to this PR after it's done

@penpornk penpornk self-assigned this Feb 1, 2024
@penpornk penpornk added the kokoro:force-run Forces CI to rerun label Feb 8, 2024
Copy link
Member

@penpornk penpornk left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the PR! Could you please help make it buildable when SYCL is not configured? (All presubmit CIs are failing. I'm not sure why all their error logs aren't viewable though.)

The Intel Extension for OpenXLA plug-in requires oneAPI Base Toolkit installation for DPC++, oneMKL, and oneTBB. I don't see their headers in this PR. Does this mean we can build this PR with SYCL without installing the oneAPI Base Toolkit?

#include "xla/stream_executor/platform/port.h"
#include "xla/stream_executor/sycl/sycl_gpu_runtime.h"

#include <level_zero/ze_api.h>
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does the user need to install Level Zero driver (intel-level-zero-gpu and level-zero) to build this file, or is the headers package (level-zero-dev) enough?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

level-zero-dev is enough for building.

@kokoro-team kokoro-team removed the kokoro:force-run Forces CI to rerun label Feb 18, 2024
@ShengYang1
Copy link
Contributor Author

ShengYang1 commented Feb 18, 2024

Thank you for the PR! Could you please help make it buildable when SYCL is not configured? (All presubmit CIs are failing. I'm not sure why all their error logs aren't viewable though.)

Sorry for the late reply. I have rebased the code and now it should work.

The Intel Extension for OpenXLA plug-in requires oneAPI Base Toolkit installation for DPC++, oneMKL, and oneTBB. I don't see their headers in this PR. Does this mean we can build this PR with SYCL without installing the oneAPI Base Toolkit?

oneAPI Base Toolkit is needed for the headers, for example sycl.hpp is included here.

@penpornk penpornk added the kokoro:force-run Forces CI to rerun label Feb 19, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Forces CI to rerun label Feb 19, 2024
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Mar 12, 2024
Imported from GitHub PR openxla/xla#10244

It is a sub PR of openxla/xla#9042.
Then we can build xla with sycl support like: `bazel build --config=sycl xxx`
Copybara import of the project:

--
e1bc7ac90f3403a477ba14533dea972de6a2c262 by Sheng, Yang <[email protected]>:

Add sycl build target

Merging this change closes #10244

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#10244 from Intel-tensorflow:yang/syclbuild e1bc7ac90f3403a477ba14533dea972de6a2c262
PiperOrigin-RevId: 615176784
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Mar 15, 2024
Imported from GitHub PR openxla/xla#10247

It is a sub PR of openxla/xla#9042.

Copybara import of the project:

--
c2f60c318aebfe4e64bf3becfd40d04949c50b41 by Sheng, Yang <[email protected]>:

Add sycl name in se_gpu_pjrt_client

Merging this change closes #10247

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#10247 from Intel-tensorflow:yang/pjrt c2f60c318aebfe4e64bf3becfd40d04949c50b41
PiperOrigin-RevId: 615176920
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Mar 15, 2024
Imported from GitHub PR openxla/xla#10247

It is a sub PR of openxla/xla#9042.

Copybara import of the project:

--
c2f60c318aebfe4e64bf3becfd40d04949c50b41 by Sheng, Yang <[email protected]>:

Add sycl name in se_gpu_pjrt_client

Merging this change closes #10247

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#10247 from Intel-tensorflow:yang/pjrt c2f60c318aebfe4e64bf3becfd40d04949c50b41
PiperOrigin-RevId: 615176920
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Mar 15, 2024
Imported from GitHub PR openxla/xla#10247

It is a sub PR of openxla/xla#9042.

Copybara import of the project:

--
c2f60c318aebfe4e64bf3becfd40d04949c50b41 by Sheng, Yang <[email protected]>:

Add sycl name in se_gpu_pjrt_client

Merging this change closes #10247

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#10247 from Intel-tensorflow:yang/pjrt c2f60c318aebfe4e64bf3becfd40d04949c50b41
PiperOrigin-RevId: 615176920
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Mar 18, 2024
Imported from GitHub PR openxla/xla#10247

It is a sub PR of openxla/xla#9042.

Copybara import of the project:

--
c2f60c318aebfe4e64bf3becfd40d04949c50b41 by Sheng, Yang <[email protected]>:

Add sycl name in se_gpu_pjrt_client

Merging this change closes #10247

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#10247 from Intel-tensorflow:yang/pjrt c2f60c318aebfe4e64bf3becfd40d04949c50b41
PiperOrigin-RevId: 615176920
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 13, 2024
Imported from GitHub PR openxla/xla#11425

It is a sub PR of openxla/xla#9042 to add sycl platform
Copybara import of the project:

--
7d0c37057c673f784089ce961d24276a496d43bd by Sheng, Yang <[email protected]>:

Add sycl platform

--
8d1196831a20f69ba8880d8549007358928f6be1 by Sheng, Yang <[email protected]>:

Add comments back

Merging this change closes #11425

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11425 from Intel-tensorflow:yang/sycl-platform 8d1196831a20f69ba8880d8549007358928f6be1
PiperOrigin-RevId: 634288679
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 13, 2024
Imported from GitHub PR openxla/xla#11425

It is a sub PR of openxla/xla#9042 to add sycl platform
Copybara import of the project:

--
7d0c37057c673f784089ce961d24276a496d43bd by Sheng, Yang <[email protected]>:

Add sycl platform

--
8d1196831a20f69ba8880d8549007358928f6be1 by Sheng, Yang <[email protected]>:

Add comments back

Merging this change closes #11425

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11425 from Intel-tensorflow:yang/sycl-platform 8d1196831a20f69ba8880d8549007358928f6be1
PiperOrigin-RevId: 634288679
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 13, 2024
Imported from GitHub PR openxla/xla#11425

It is a sub PR of openxla/xla#9042 to add sycl platform
Copybara import of the project:

--
7d0c37057c673f784089ce961d24276a496d43bd by Sheng, Yang <[email protected]>:

Add sycl platform

--
8d1196831a20f69ba8880d8549007358928f6be1 by Sheng, Yang <[email protected]>:

Add comments back

Merging this change closes #11425

PiperOrigin-RevId: 652040832
copybara-service bot pushed a commit to google/tsl that referenced this pull request Aug 8, 2024
Imported from GitHub PR openxla/xla#11424

It is a sub PR of openxla/xla#9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76e2b71120106ffae91945df9b974e74dec by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76e2b71120106ffae91945df9b974e74dec
PiperOrigin-RevId: 660816063
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Aug 8, 2024
Imported from GitHub PR openxla/xla#11424

It is a sub PR of openxla/xla#9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76e2b71120106ffae91945df9b974e74dec by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76e2b71120106ffae91945df9b974e74dec
PiperOrigin-RevId: 660816063
copybara-service bot pushed a commit to google/tsl that referenced this pull request Aug 8, 2024
Imported from GitHub PR openxla/xla#11424

It is a sub PR of openxla/xla#9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76e2b71120106ffae91945df9b974e74dec by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76e2b71120106ffae91945df9b974e74dec
PiperOrigin-RevId: 660816063
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Aug 8, 2024
Imported from GitHub PR openxla/xla#11424

It is a sub PR of openxla/xla#9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76e2b71120106ffae91945df9b974e74dec by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76e2b71120106ffae91945df9b974e74dec
PiperOrigin-RevId: 660816063
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Aug 8, 2024
Imported from GitHub PR openxla/xla#11424

It is a sub PR of openxla/xla#9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76e2b71120106ffae91945df9b974e74dec by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76e2b71120106ffae91945df9b974e74dec
PiperOrigin-RevId: 660816063
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Aug 8, 2024
Imported from GitHub PR openxla/xla#11424

It is a sub PR of openxla/xla#9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76e2b71120106ffae91945df9b974e74dec by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76e2b71120106ffae91945df9b974e74dec
PiperOrigin-RevId: 660816063
copybara-service bot pushed a commit to google/tsl that referenced this pull request Aug 9, 2024
Imported from GitHub PR openxla/xla#11424

It is a sub PR of openxla/xla#9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76e2b71120106ffae91945df9b974e74dec by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76e2b71120106ffae91945df9b974e74dec
PiperOrigin-RevId: 660816063
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Aug 9, 2024
Imported from GitHub PR openxla/xla#11424

It is a sub PR of openxla/xla#9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76e2b71120106ffae91945df9b974e74dec by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76e2b71120106ffae91945df9b974e74dec
PiperOrigin-RevId: 660816063
copybara-service bot pushed a commit that referenced this pull request Aug 9, 2024
Imported from GitHub PR #11424

It is a sub PR of #9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76 by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

FUTURE_COPYBARA_INTEGRATE_REVIEW=#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76
PiperOrigin-RevId: 661176906
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Aug 9, 2024
Imported from GitHub PR openxla/xla#11424

It is a sub PR of openxla/xla#9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76e2b71120106ffae91945df9b974e74dec by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76e2b71120106ffae91945df9b974e74dec
PiperOrigin-RevId: 661176906
copybara-service bot pushed a commit that referenced this pull request Aug 9, 2024
Imported from GitHub PR #11424

It is a sub PR of #9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76 by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

FUTURE_COPYBARA_INTEGRATE_REVIEW=#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76
PiperOrigin-RevId: 661176906
copybara-service bot pushed a commit that referenced this pull request Aug 9, 2024
Imported from GitHub PR #11424

It is a sub PR of #9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76 by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

COPYBARA_INTEGRATE_REVIEW=#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76
PiperOrigin-RevId: 661186664
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Aug 9, 2024
Imported from GitHub PR openxla/xla#11424

It is a sub PR of openxla/xla#9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76e2b71120106ffae91945df9b974e74dec by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

PiperOrigin-RevId: 661186664
copybara-service bot pushed a commit to google/tsl that referenced this pull request Aug 9, 2024
Imported from GitHub PR openxla/xla#11424

It is a sub PR of openxla/xla#9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76e2b71120106ffae91945df9b974e74dec by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76e2b71120106ffae91945df9b974e74dec
PiperOrigin-RevId: 661222386
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Aug 9, 2024
Imported from GitHub PR openxla/xla#11424

It is a sub PR of openxla/xla#9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76e2b71120106ffae91945df9b974e74dec by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76e2b71120106ffae91945df9b974e74dec
PiperOrigin-RevId: 661222386
copybara-service bot pushed a commit to google/tsl that referenced this pull request Aug 9, 2024
Imported from GitHub PR openxla/xla#11424

It is a sub PR of openxla/xla#9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76e2b71120106ffae91945df9b974e74dec by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

PiperOrigin-RevId: 661272459
copybara-service bot pushed a commit that referenced this pull request Aug 9, 2024
Imported from GitHub PR #11424

It is a sub PR of #9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76 by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

COPYBARA_INTEGRATE_REVIEW=#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76
PiperOrigin-RevId: 661272459
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Aug 9, 2024
Imported from GitHub PR openxla/xla#11424

It is a sub PR of openxla/xla#9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76e2b71120106ffae91945df9b974e74dec by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

PiperOrigin-RevId: 661272459
sergachev pushed a commit that referenced this pull request Aug 12, 2024
Imported from GitHub PR #11424

It is a sub PR of #9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76 by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

COPYBARA_INTEGRATE_REVIEW=#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76
PiperOrigin-RevId: 661186664
sergachev pushed a commit that referenced this pull request Aug 12, 2024
Imported from GitHub PR #11424

It is a sub PR of #9042 to add spirv-llvm-translator and translation pass
Copybara import of the project:

--
6d8ce76 by Sheng, Yang <[email protected]>:

Add SPIRV-LLVM-Translator and translation pass

update comments and SPIRV-LLVM-Translator commit

Merging this change closes #11424

COPYBARA_INTEGRATE_REVIEW=#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76
PiperOrigin-RevId: 661272459
steeve pushed a commit to zml/xla that referenced this pull request Aug 30, 2024
Imported from GitHub PR openxla#10247

It is a sub PR of openxla#9042.

Copybara import of the project:

--
c2f60c3 by Sheng, Yang <[email protected]>:

Add sycl name in se_gpu_pjrt_client

Merging this change closes openxla#10247

COPYBARA_INTEGRATE_REVIEW=openxla#10247 from Intel-tensorflow:yang/pjrt c2f60c3
PiperOrigin-RevId: 617141542
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.

7 participants