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.

❌ 4 New Failures, 1 Pending

As of commit 2272a3f with merge base 9bcd73b ():

NEW FAILURES - The following jobs have failed:

• Build M1 Wheels / pytorch/ao / upload / wheel-py3_9-cpu (gh)
  The process '/usr/bin/git' failed with exit code 1
• Code Analysis with Ruff / build (3.9) (gh)
  Process completed with exit code 1.
• Run TorchAO Experimental Tests / test-cpu-ops (macos-14) (gh)
  torchao/experimental/tests/test_int8_dynamic_activation_intx_weight.py::TestInt8DynamicActivationIntxWeight::test_export_compile_aoti_PackedLinearInt8DynamicActivationIntxWeightLayout
• Run TorchAO Experimental Tests / test-mps-ops (macos-m1-stable) (gh)
  AttributeError: '_OpNamespace' 'torchao' object has no attribute '_linear_fp_act_1bit_weight'

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

[jcaip]

This looks good, but there are no tests? We should be able to turn on the int4wo tests to run on ROCm to verify this is correct right? Or is there additional work to be done to turn on the test?

[petrex]

This looks good, but there are no tests? We should be able to turn on the int4wo tests to run on ROCm to verify this is correct right? Or is there additional work to be done to turn on the test?

Thanks , This PR introduced optimized kernel for MI300x. We are working on allocating Mi300x nodes for torchao CI.
Will enable mi300x specific tests in following PRs.

[jcaip]

@petrex OK, that makes sense, approving for now.

We should also run some updated benchmarks now that we have int4+ marlin support. I'll try to find some time next week to do so.

[jcaip]

cc @petrex I'm assuming you want to merge in this first and then merge #1206 in? In that case I think it should be okay so long as CI passes for 1206 since you already have approval on that PR.