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

.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

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:

src/dynamics_kernel.cu

@@ -292,7 +292,7 @@ std::vector<torch::Tensor> 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<scalar_t><<<blocks, threads>>>(
             R.packed_accessor<scalar_t,3,torch::RestrictPtrTraits,size_t>(),
             dg.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
@@ -338,7 +338,7 @@ std::vector<torch::Tensor> 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<scalar_t><<<blocks, threads>>>(
             R.packed_accessor<scalar_t,3,torch::RestrictPtrTraits,size_t>(),
             dg.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
@@ -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<scalar_t><<<blocks, threads>>>(
             R_new.packed_accessor<scalar_t,3,torch::RestrictPtrTraits,size_t>(),
             R.packed_accessor<scalar_t,3,torch::RestrictPtrTraits,size_t>(),

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<scalar_t><<<blocks, threads>>>(
             canvas.packed_accessor<scalar_t,3,torch::RestrictPtrTraits,size_t>(),
             flow.packed_accessor<scalar_t,4,torch::RestrictPtrTraits,size_t>(),
@@ -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<scalar_t><<<blocks, threads>>>(
             depth.packed_accessor<scalar_t,4,torch::RestrictPtrTraits,size_t>(),
             dddp.packed_accessor<scalar_t,4,torch::RestrictPtrTraits,size_t>(),
@@ -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<scalar_t><<<blocks, threads>>>(
             nearest_pt.packed_accessor<scalar_t,3,torch::RestrictPtrTraits,size_t>(),
             balls.packed_accessor<scalar_t,3,torch::RestrictPtrTraits,size_t>(),

src/setup.py

@@ -1,14 +1,43 @@
+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
+    # <torch/extension.h> (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=[
         CUDAExtension('quadsim_cuda', [
             'quadsim.cpp',
             'quadsim_kernel.cu',
             'dynamics_kernel.cu',
-        ]),
+        ], extra_link_args=extra_link_args),
     ],
     cmdclass={
         'build_ext': BuildExtension