Build the quadsim_cuda extension on AMD GPUs (ROCm/HIP)#45
Open
jeffdaily wants to merge 1 commit into
Open
Conversation
Add AMD GPU support to the quadsim_cuda PyTorch extension. PyTorch's
build-time hipify translates the CUDA sources to HIP, so the same
`pip install -e src` builds for ROCm with no source restructuring.
Review order:
- src/dynamics_kernel.cu, src/quadsim_kernel.cu: at the six
AT_DISPATCH_FLOATING_TYPES sites, replace the deprecated
Tensor.type() with Tensor.scalar_type(). Tensor.type() returns
DeprecatedTypeProperties, which recent PyTorch no longer converts to
ScalarType, so the dispatch macro fails to compile. This is a PyTorch
forward-compatibility fix and is identical on CUDA; the dispatched
scalar_t is unchanged, so it is value-preserving.
- src/setup.py: on Windows, c10.dll does not export the
c10::ValueError(SourceLocation, string) constructor that headers
pulled in via <torch/extension.h> reference, causing LNK2001. Alias it
to the exported c10::Error(SourceLocation, string) constructor with a
linker /ALTERNATENAME; ValueError IS-A Error with no extra members, so
the constructors are layout- and semantics-identical. Guarded to
win32, so Linux and CUDA builds are untouched.
- .gitignore: ignore Python and torch-extension build outputs (the
hipify .hip mirrors, the compiled .so, build/, egg-info).
- README.md: note in the existing build section that the ops also build
on AMD GPUs with a ROCm build of PyTorch.
The kernels are all one-thread-per-output with no warp-width-sensitive
constructs (no shuffles, ballots, shared memory, or cub), so the wave64
AMD targets produce identical results to wave32 and CUDA.
Test Plan:
Built and ran the repo's own test.py (forward-sim allclose plus
analytic-backward vs torch.autograd on 64 random double states) on AMD
GPUs.
```
export HIP_VISIBLE_DEVICES=0 PYTORCH_ROCM_ARCH=gfx90a
cd src && rm -rf build *.egg-info && pip install -e . --no-build-isolation
python3 -c "import torch; torch.manual_seed(42)
import runpy; runpy.run_path('test.py', run_name='__main__')"
```
All four forward outputs and all five gradients pass torch.allclose.
Validated on gfx90a (MI250X) and gfx1100 (Radeon Pro W7800) on Linux
ROCm 7.2, and on Windows ROCm (gfx1101, gfx1201). Native gfx code object
dispatch confirmed via AMD_LOG_LEVEL=3 (no host fallback).
test.py as shipped uses unseeded randoms and its v_next allclose flaps
occasionally; this is a pre-existing upstream numeric artifact (the
kernel takes ctl_dt as float while the Python reference uses a double
1/15) that reproduces on stock CUDA, not a port regression. Run with a
fixed seed for a deterministic gate.
Authored with assistance from Claude (Anthropic).
This file contains hidden or 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
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 adds AMD GPU support to the
quadsim_cudaPyTorch extension. PyTorch's build-time hipify translates the CUDA sources to HIP, sopip install -e srcbuilds for ROCm with no source restructuring.Changes:
dynamics_kernel.cu,quadsim_kernel.cu: at the sixAT_DISPATCH_FLOATING_TYPESsites, useTensor.scalar_type()instead of the deprecatedTensor.type().Tensor.type()returnsDeprecatedTypeProperties, which recent PyTorch no longer converts toScalarType, so the dispatch macro fails to compile. This is a PyTorch forward-compatibility fix that applies identically on CUDA; the dispatchedscalar_tis unchanged, so it is value-preserving.setup.py: on Windows built with clang-cl,c10.dlldoes not export thec10::ValueError(SourceLocation, string)constructor inherited viausing Error::Error(clang-cl omits dllexport for inherited constructors, [clang-cl] Inherited constructors dll export discrepancy compared to MSVC llvm/llvm-project#162640), so headers reaching it through<torch/extension.h>fail to link with LNK2001. Alias it to the exportedc10::Error(SourceLocation, string)constructor via a linker/ALTERNATENAME;ValueErrorIS-AErrorwith no extra members, so they are layout- and semantics-identical. This is fixed upstream in [c10] Fix missing symbol exports for ValueError/NotImplementedError on Windows pytorch/pytorch#175340 -- the alias is a fallback that goes unused once PyTorch includes that fix, and is guarded to win32, so Linux and CUDA builds are untouched..gitignore: ignore Python and torch-extension build outputs.README.md: note in the build section that the ops also build on AMD GPUs.The kernels are all one-thread-per-output with no warp-width-sensitive constructs, so the wave64 AMD targets produce results identical to CUDA.
Testing
Built and ran the repo's own
test.py(forward-sim allclose plus analytic-backward vstorch.autogradon 64 random double states):All four forward outputs and all five gradients pass
torch.allclose. Validated on gfx90a (MI250X) and gfx1100 (Radeon Pro W7800) on Linux ROCm 7.2, and on Windows ROCm (gfx1201, gfx1101). Native gfx code-object dispatch confirmed viaAMD_LOG_LEVEL=3(no host fallback).test.pyas shipped uses unseeded randoms and itsv_nextallclose flaps occasionally; this is a pre-existing upstream numeric artifact (the kernel takesctl_dtasfloatwhile the Python reference uses a double1/15) that reproduces on stock CUDA, not a regression here. Run with a fixed seed for a deterministic gate.This port was authored with assistance from Claude (Anthropic).