ROCm Support : Tile_Layout kernel

[petrex]

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_layout as proof of concept, but it generalizes to other kernels (for example, fp6_llm or sparse_marlin) with minimum effort. Feedback are welcome

co-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 install

Next

• AMD specific unit tests (tensor_core_tiled_layout)
• workload and platform specific optimization (tensor_core_tiled_layout)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1201

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 89bae5f with merge base e1cb44a ():

NEW FAILURE - The following job has failed:

• Build M1 Wheels / pytorch/ao / upload / wheel-py3_9-cpu (gh)
  The process '/usr/bin/git' failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[msaroufim]

Thanks! Left some prelim feedback more of a n00b question but is this the main way you've seen ROCm supported in new codebases? The other approach would be to have some code duplication and avoid a boolean flag like USE_ROCM in the .cu file

[msaroufim]

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

[petrex]

thanks for approving the CI flow, we will look into this

[petrex]

@msaroufim can we have the CI flow to test the change? thanks

[msaroufim]

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

[petrex]

@msaroufim happy new year! The 2 failing CI check seems irrelevant to this PR? Any tips?

  /usr/bin/git -c protocol.version=2 fetch --no-tags --prune --no-recurse-submodules --depth=1 origin +refs/heads/rocm_enablement_staging*:refs/remotes/origin/rocm_enablement_staging* +refs/tags/rocm_enablement_staging*:refs/tags/rocm_enablement_staging*
  Error: The process '/usr/bin/git' failed with exit code 1

[jerryzh168]

thanks! this is the dequantize and unpack_dequantize op btw, the packing op and int4mm op lives in pytorch right now: https://github.com/pytorch/pytorch/blob/b1c2c3967adf035356abfcb90d0b39ff7ee63707/aten/src/ATen/native/native_functions.yaml#L4162-L4170

but this can be used to test build for ROCm