From c018c93b2c008c09eed798243c24edee10be5251 Mon Sep 17 00:00:00 2001 From: Jeff Daily Date: Thu, 25 Jun 2026 22:23:40 +0000 Subject: [PATCH] [ROCm] Build quadsim_cuda extension on ROCm/HIP 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 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). --- .gitignore | 9 +++++++++ README.md | 2 ++ src/dynamics_kernel.cu | 6 +++--- src/quadsim_kernel.cu | 6 +++--- src/setup.py | 31 ++++++++++++++++++++++++++++++- 5 files changed, 47 insertions(+), 7 deletions(-) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ff13f8b --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +# Python / build artifacts +__pycache__/ +*.py[cod] +*.egg-info/ +build/ + +# torch CUDAExtension build outputs (hipify mirrors and compiled module) +*.hip +*.so diff --git a/README.md b/README.md index 55998ba..d3aa48b 100644 --- a/README.md +++ b/README.md @@ -43,6 +43,8 @@ To build the CUDA operations, run the following command: pip install -e src ``` +The same command also builds the operations on AMD GPUs: with a [ROCm](https://rocm.docs.amd.com/) build of PyTorch it works unchanged, since PyTorch translates the CUDA sources to HIP at build time. Optionally set `PYTORCH_ROCM_ARCH` (for example `gfx90a` or `gfx1100`) to build only for your GPU. Tested with ROCm 7.2 and PyTorch 2.13 on gfx90a (MI250X) and gfx1100 (Radeon Pro W7800). + ## Training To start the training process, use the following command: diff --git a/src/dynamics_kernel.cu b/src/dynamics_kernel.cu index cd9a026..ae3153e 100644 --- a/src/dynamics_kernel.cu +++ b/src/dynamics_kernel.cu @@ -292,7 +292,7 @@ std::vector run_forward_cuda( const int threads = R.size(0); const dim3 blocks(1); - AT_DISPATCH_FLOATING_TYPES(R.type(), "run_forward_cuda", ([&] { + AT_DISPATCH_FLOATING_TYPES(R.scalar_type(), "run_forward_cuda", ([&] { run_forward_cuda_kernel<<>>( R.packed_accessor(), dg.packed_accessor(), @@ -338,7 +338,7 @@ std::vector run_backward_cuda( const int threads = R.size(0); const dim3 blocks(1); - AT_DISPATCH_FLOATING_TYPES(R.type(), "run_backward_cuda", ([&] { + AT_DISPATCH_FLOATING_TYPES(R.scalar_type(), "run_backward_cuda", ([&] { run_backward_cuda_kernel<<>>( R.packed_accessor(), dg.packed_accessor(), @@ -371,7 +371,7 @@ torch::Tensor update_state_vec_cuda( const int threads = a_thr.size(0); const dim3 blocks(1); torch::Tensor R_new = torch::empty_like(R); - AT_DISPATCH_FLOATING_TYPES(a_thr.type(), "update_state_vec", ([&] { + AT_DISPATCH_FLOATING_TYPES(a_thr.scalar_type(), "update_state_vec", ([&] { update_state_vec_cuda_kernel<<>>( R_new.packed_accessor(), R.packed_accessor(), diff --git a/src/quadsim_kernel.cu b/src/quadsim_kernel.cu index 434056d..0fdbd97 100644 --- a/src/quadsim_kernel.cu +++ b/src/quadsim_kernel.cu @@ -333,7 +333,7 @@ void render_cuda( size_t state_size = canvas.numel(); const dim3 blocks((state_size + threads - 1) / threads); - AT_DISPATCH_FLOATING_TYPES(canvas.type(), "render_cuda", ([&] { + AT_DISPATCH_FLOATING_TYPES(canvas.scalar_type(), "render_cuda", ([&] { render_cuda_kernel<<>>( canvas.packed_accessor(), flow.packed_accessor(), @@ -359,7 +359,7 @@ void rerender_backward_cuda( size_t state_size = dddp.numel(); const dim3 blocks((state_size + threads - 1) / threads); - AT_DISPATCH_FLOATING_TYPES(depth.type(), "rerender_backward_cuda", ([&] { + AT_DISPATCH_FLOATING_TYPES(depth.scalar_type(), "rerender_backward_cuda", ([&] { rerender_backward_cuda_kernel<<>>( depth.packed_accessor(), dddp.packed_accessor(), @@ -379,7 +379,7 @@ void find_nearest_pt_cuda( const int threads = 1024; size_t state_size = pos.size(0) * pos.size(1); const dim3 blocks((state_size + threads - 1) / threads); - AT_DISPATCH_FLOATING_TYPES(pos.type(), "nearest_pt_cuda", ([&] { + AT_DISPATCH_FLOATING_TYPES(pos.scalar_type(), "nearest_pt_cuda", ([&] { nearest_pt_cuda_kernel<<>>( nearest_pt.packed_accessor(), balls.packed_accessor(), diff --git a/src/setup.py b/src/setup.py index b08e682..e977776 100644 --- a/src/setup.py +++ b/src/setup.py @@ -1,6 +1,35 @@ +import sys + from setuptools import setup from torch.utils.cpp_extension import BuildExtension, CUDAExtension +extra_link_args = [] +if sys.platform == "win32": + # On PyTorch 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 (llvm/llvm-project#162640). Headers pulled in by + # (e.g. ATen/TensorIndexing.h) use TORCH_CHECK_VALUE, + # which emits a __declspec(dllimport) reference to that constructor, so the + # link fails with LNK2001. Alias it to Error(SourceLocation, string), which + # IS exported; ValueError IS-A Error with no extra data members, so the two + # constructors are layout- and semantics-identical. + # + # Fixed upstream in https://github.com/pytorch/pytorch/pull/175340, which + # adds the explicit export. /ALTERNATENAME only supplies a fallback when + # the symbol is otherwise unresolved, so on a PyTorch that includes that fix + # the real export resolves and this alias is silently unused -- it is needed + # only for older builds that predate it. + _val_imp = ( + "__imp_??0ValueError@c10@@QEAA@USourceLocation@1@" + "V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z" + ) + _err_imp = ( + "__imp_??0Error@c10@@QEAA@USourceLocation@1@" + "V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z" + ) + extra_link_args.append(f"/ALTERNATENAME:{_val_imp}={_err_imp}") + setup( name='quadsim_cuda', ext_modules=[ @@ -8,7 +37,7 @@ 'quadsim.cpp', 'quadsim_kernel.cu', 'dynamics_kernel.cu', - ]), + ], extra_link_args=extra_link_args), ], cmdclass={ 'build_ext': BuildExtension