Skip to content

Build the quadsim_cuda extension on AMD GPUs (ROCm/HIP)#45

Open
jeffdaily wants to merge 1 commit into
HenryHuYu:masterfrom
jeffdaily:moat-port
Open

Build the quadsim_cuda extension on AMD GPUs (ROCm/HIP)#45
jeffdaily wants to merge 1 commit into
HenryHuYu:masterfrom
jeffdaily:moat-port

Conversation

@jeffdaily

Copy link
Copy Markdown

This adds AMD GPU support to the quadsim_cuda PyTorch extension. PyTorch's build-time hipify translates the CUDA sources to HIP, so pip install -e src builds for ROCm with no source restructuring.

Changes:

  • dynamics_kernel.cu, quadsim_kernel.cu: at the six AT_DISPATCH_FLOATING_TYPES sites, use Tensor.scalar_type() instead of the deprecated Tensor.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 that applies identically on CUDA; the dispatched scalar_t is unchanged, so it is value-preserving.
  • setup.py: on Windows built with clang-cl, c10.dll does not export the c10::ValueError(SourceLocation, string) constructor inherited via using 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 exported c10::Error(SourceLocation, string) constructor via a linker /ALTERNATENAME; ValueError IS-A Error with 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 vs torch.autograd on 64 random double states):

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 (gfx1201, gfx1101). 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 regression here. Run with a fixed seed for a deterministic gate.

This port was authored with assistance from Claude (Anthropic).

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).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant