Support lowering XLA clamp instruction to cuDNN. #13020
Closed
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
sergeykozub
approved these changes
May 24, 2024
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
May 24, 2024
Imported from GitHub PR openxla/xla#13020 Support lowering XLA clamp instruction to cuDNN. cc @sergachev Copybara import of the project: -- 47dc71f2a0d5887461a0b7d985328442e0e8da2f by Elfie Guo <elfieg@nvidia.com>: Support lowering clamp instruction to cuDNN. Merging this change closes #13020 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#13020 from elfiegg:fp8_triton 47dc71f2a0d5887461a0b7d985328442e0e8da2f PiperOrigin-RevId: 636814740
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
May 24, 2024
Imported from GitHub PR openxla/xla#13020 Support lowering XLA clamp instruction to cuDNN. cc @sergachev Copybara import of the project: -- 47dc71f2a0d5887461a0b7d985328442e0e8da2f by Elfie Guo <elfieg@nvidia.com>: Support lowering clamp instruction to cuDNN. Merging this change closes #13020 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#13020 from elfiegg:fp8_triton 47dc71f2a0d5887461a0b7d985328442e0e8da2f PiperOrigin-RevId: 636814740
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
May 24, 2024
Imported from GitHub PR openxla/xla#13020 Support lowering XLA clamp instruction to cuDNN. cc @sergachev Copybara import of the project: -- 47dc71f2a0d5887461a0b7d985328442e0e8da2f by Elfie Guo <elfieg@nvidia.com>: Support lowering clamp instruction to cuDNN. Merging this change closes #13020 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#13020 from elfiegg:fp8_triton 47dc71f2a0d5887461a0b7d985328442e0e8da2f PiperOrigin-RevId: 636814740
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
May 24, 2024
Imported from GitHub PR openxla/xla#13020 Support lowering XLA clamp instruction to cuDNN. cc @sergachev Copybara import of the project: -- 47dc71f2a0d5887461a0b7d985328442e0e8da2f by Elfie Guo <elfieg@nvidia.com>: Support lowering clamp instruction to cuDNN. Merging this change closes #13020 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#13020 from elfiegg:fp8_triton 47dc71f2a0d5887461a0b7d985328442e0e8da2f PiperOrigin-RevId: 636814740
sergeykozub
approved these changes
May 24, 2024
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
May 24, 2024
Imported from GitHub PR openxla/xla#13020 Support lowering XLA clamp instruction to cuDNN. cc @sergachev Copybara import of the project: -- 47dc71f2a0d5887461a0b7d985328442e0e8da2f by Elfie Guo <elfieg@nvidia.com>: Support lowering clamp instruction to cuDNN. Merging this change closes #13020 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#13020 from elfiegg:fp8_triton 47dc71f2a0d5887461a0b7d985328442e0e8da2f PiperOrigin-RevId: 636814740
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
May 24, 2024
Imported from GitHub PR openxla/xla#13020 Support lowering XLA clamp instruction to cuDNN. cc @sergachev Copybara import of the project: -- 47dc71f2a0d5887461a0b7d985328442e0e8da2f by Elfie Guo <elfieg@nvidia.com>: Support lowering clamp instruction to cuDNN. Merging this change closes #13020 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#13020 from elfiegg:fp8_triton 47dc71f2a0d5887461a0b7d985328442e0e8da2f PiperOrigin-RevId: 636814740
copybara-service bot
pushed a commit
that referenced
this pull request
May 24, 2024
Overall the idea is to collect profile data for each module given amount of times (which can be configured) then recompile the module with the aggregated profile data. 1. We need to track how many times each module were profiled and collect profiling results. For this i added a ProfileSessionRunner class at profile.py. The class can track how many times an instance of it was called to profile a session and also can aggregate profile results. 2. We need associate profiling session to the module at the interpreter. To do this i added a dictionary to pjit.py which associates Jaxpr with profile session runner. 3. The profile session runner should be passed to pxla.py and then called. 4. We need to correctly deal with fast path at the interpreter level, so JAX won't use HLO directly if PGLE need to be collected, but also JAX will not recompiled the module only for PGLE. See changes in pjit.py and in lru_cache.h 5. Once FDO is collected we need to share it between hosts to keep deterministic compilation. FUTURE_COPYBARA_INTEGRATE_REVIEW=#13020 from elfiegg:fp8_triton 47dc71f PiperOrigin-RevId: 617824989
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
May 24, 2024
Imported from GitHub PR openxla/xla#13020 Support lowering XLA clamp instruction to cuDNN. cc @sergachev Copybara import of the project: -- 47dc71f2a0d5887461a0b7d985328442e0e8da2f by Elfie Guo <elfieg@nvidia.com>: Support lowering clamp instruction to cuDNN. Merging this change closes #13020 PiperOrigin-RevId: 636875273
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
May 24, 2024
Overall the idea is to collect profile data for each module given amount of times (which can be configured) then recompile the module with the aggregated profile data. 1. We need to track how many times each module were profiled and collect profiling results. For this i added a ProfileSessionRunner class at profile.py. The class can track how many times an instance of it was called to profile a session and also can aggregate profile results. 2. We need associate profiling session to the module at the interpreter. To do this i added a dictionary to pjit.py which associates Jaxpr with profile session runner. 3. The profile session runner should be passed to pxla.py and then called. 4. We need to correctly deal with fast path at the interpreter level, so JAX won't use HLO directly if PGLE need to be collected, but also JAX will not recompiled the module only for PGLE. See changes in pjit.py and in lru_cache.h 5. Once FDO is collected we need to share it between hosts to keep deterministic compilation. FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#13020 from elfiegg:fp8_triton 47dc71f2a0d5887461a0b7d985328442e0e8da2f PiperOrigin-RevId: 617824989
|
This PR results in "Unsupported elementwise operation %clamp.1" in an internal test, which is raised in "ir_emitter_triton.cc" |
|
This PR was rolled back in d114ece! |
|
sent https://github.com/openxla/xla/pull/13061/files, can you help verify if this fixes the test? |
|
Confirmed, this fixed the test, and we have rolled this PR forward again. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Support lowering XLA clamp instruction to cuDNN.
cc @sergachev