ROCm Support : HIP kernel Generator #1201
Open
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
ROCm build infrastructure
[ROCm] Enable Tiled layout extension and minor changes to setup
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1201
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ❌ 10 New Failures, 2 Unrelated FailuresAs of commit 612ad14 with merge base 6234116 ( NEW FAILURES - The following jobs have failed:
BROKEN TRUNK - The following jobs failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
msaroufim
reviewed
Nov 4, 2024
msaroufim
reviewed
Nov 4, 2024
| @@ -57,11 +59,11 @@ def get_extensions(): | |||
|
|
|||
| if not torch.cuda.is_available(): | |||
| print("PyTorch GPU support is not available. Skipping compilation of CUDA extensions") | |||
| if CUDA_HOME is None and torch.cuda.is_available(): | |||
| print("CUDA toolkit is not available. Skipping compilation of CUDA extensions") | |||
| if CUDA_HOME is None or not IS_ROCM and torch.cuda.is_available(): | |||
There was a problem hiding this comment.
some of these if conditions might cause our binary builds to break - check CI, since we run binary builds when people modify setup.py
Until you get your first PR merged I'll have to manually approve every CI run so might be easiest if you wanna make some simple doc level PR and I can merge that fast
There was a problem hiding this comment.
thanks for approving the CI flow, we will look into this
|
@msaroufim can we have the CI flow to test the change? thanks |
|
I gave you triage access @petrex - what might help is also merging some simple PR since after that you shouldn't need me to trigger CI for you |
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.
This pull request introduces support for ROCm (Radeon Open Compute) in addition to CUDA for GPU acceleration. The changes primarily focus on enabling the build and execution of ROCm-specific code paths alongside existing CUDA paths.
In this PR, I use
tensor_core_tiled_layoutas proof of concept, but it generalizes to other kernels (for example, fp6_llm or sparse_marlin) with minimum effort. Feedback are welcomeco-author : @lcskrishna
Features:
ROCm Support Integration:
setup.py: Added detection for ROCm and adjusted the logic for compiling GPU extensions based on the availability of CUDA or ROCm.Conditional Compilation for ROCm:
torchao/csrc/cuda/tensor_core_tiled_layout/tensor_core_tiled_layout.cu: Introduced conditional compilation directives to include ROCm-specific headers and adjust constants and operations for ROCm.These changes ensure that the codebase can compile and run efficiently on both CUDA and ROCm platforms, leveraging the best available GPU acceleration technology.
Usage
With ROCm pytorch nightly docker , simply run
PYTORCH_ROCM_ARCH=gfx942 python setup.py installNext