Port pytorch3d (#2039)

Summary:
Enables building pytorch3d's `_C` extension against a ROCm-built PyTorch and running the test suite on AMD GPUs, including the pulsar subrenderer. Verified on AMD Instinct MI250X (gfx90a, warpSize=64), HIP 7.2, PyTorch 2.13.

## Mechanics

`torch.utils.cpp_extension.BuildExtension` auto-hipifies `.cu` sources of a `CUDAExtension` against a HIP-built torch (`cuda_runtime.h → hip/hip_runtime.h`, `cub:: → hipcub::`, `cudaStream_t → hipStream_t`, etc.), so most of the lift is build-system glue and a small number of CUDA intrinsics that don't have HIP equivalents.

- `setup.py`: detect ROCm via `torch.version.hip is not None`; treat `ROCM_HOME` as the GPU-toolkit-root analogue of `CUDA_HOME` (without this, `CUDA_HOME is None` silently demoted the build to a CPU-only `CppExtension`); skip `CUB_HOME`, CUDA-13 visibility flags, and `-ccbin=` on ROCm.
- `pytorch3d/csrc/pulsar/gpu/commands.h`: CUDA's `_rn`-suffixed FP rounding intrinsics (`__fadd_rn`, `__fdiv_rn`, `__fsqrt_rn`, `__fmaf_rn`, `__frcp_rn`) and `__saturatef` have no HIP equivalents — AMD's GPU ISA has no instruction-level rounding-mode override, so they expand to plain operators / `sqrtf` / `fmaf` / `1.0f/x` / `fmaxf(0,fminf(1,x))` on the `USE_ROCM` arm, which are rounding-mode-equivalent (both round-to-nearest-even). The HIP compiler may fuse `a+b*c` into a single-rounding FMA where CUDA's `_rn` would have prevented it; if FMA-fusion drift ever becomes a numerical issue, add `-ffp-contract=off` to pulsar's HIPCC flags. `__powf` is replaced with `powf`. `atomicAdd_block` has no HIP function-name equivalent — the semantic equivalent is `__hip_atomic_fetch_add(ptr, val, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_WORKGROUP)` (plain HIP `atomicAdd` is device-scope, strictly stronger than block-scope and forces L2-coherent atomics).
- `tests/test_point_mesh_distance.py`: loosen `grad_faces` tolerance in `test_point_face_distance` from `5e-7` to `5e-6` to match the sibling `test_face_point_distance`. The backward kernel uses `atomicAdd` and calls `alertNotDeterministic`; FP add order varies by wavefront width.
- The X_t / camera-R/T equality checks in `test_points_alignment.py` and `test_cameras_alignment.py` are now skipped when `n_points <= dim` (resp. `batch_size <= 3` for camera-center alignment in 3D). Mean-centering renders the SVD rank-deficient in those cases, so the rotation around the degenerate axis is non-unique and different BLAS implementations (rocBLAS RDNA vs CDNA, cuBLAS) pick different valid null-space directions. The center-alignment check still runs and verifies the well-defined part of the transformation.

Pull Request resolved: #2039

Test Plan:
All GPU tests pass on both AMD Instinct MI250X (gfx90a, wave64, HIP 7.2) and AMD Radeon Pro W7800 (gfx1100, wave32, HIP 7.2.53211, torch 2.13.0a0).

| Module | Result |
|---|---|
| knn, ball_query, sample_farthest_points, face_areas_normals | all pass |
| rasterize_points, rasterize_meshes, chamfer, packed_to_padded | all pass |
| interpolate_face_attributes, blending, compositing, sample_pdf, mesh_normal_consistency | all pass |
| point_mesh_distance | 9/9 pass (with tolerance fix in this PR) |
| pulsar/test_forward, test_channels, test_depth, test_hands, test_ortho, test_small_spheres | 10 passed (FB_TEST=1) |
| test_render_points pulsar tests, test_camera_conversions::test_pulsar_conversion | 3 passed |
| points_to_volumes, iou_box3d, marching_cubes | 20 failures, all env-only |

The 20 env-only failures are `torch.inverse()` on CPU tensors in test reference paths; this verification host's PyTorch was built with `USE_LAPACK: 0` (only `mkl-static` `.a` archives in the conda env; PyTorch's `FindBLAS` looks for `libmkl_intel_lp64.so`). Unrelated to the port — re-verifying with a LAPACK-linked PyTorch is left to upstream.

Reviewed By: MichaelRamamonjisoa

Differential Revision: D106825690

Pulled By: bottler

fbshipit-source-id: f7a9b6028e6fb555f3b8c0f9792e88b818327166

Author: jeffdaily

.github/workflows/build.yml

@@ -21,3 +21,45 @@ jobs:
       run: |-
         conda create --name env --yes --quiet conda-build
         conda run --no-capture-output --name env python3 ./packaging/build_conda.py --use-conda-cuda
+
+  # Build-only verification for the ROCm/HIP code paths. Runs in an AMD ROCm
+  # dev container on a CPU-only GitHub runner; we don't need an AMD GPU just
+  # to compile, and not running tests keeps the CI cost low. Catches build
+  # regressions in the ROCm code paths (USE_ROCM guards, hipify-touched sources,
+  # the pulsar HIP intrinsic replacements, etc.).
+  linux_rocm_build:
+    runs-on: ubuntu-latest
+    container:
+      # `-complete` tag bundles the full ROCm math stack (rocThrust, hipCUB,
+      # rocPRIM, ...). The plain `7.2.3` tag is HIP-runtime-only and fails to
+      # find <thrust/complex.h> when including PyTorch headers.
+      image: rocm/dev-ubuntu-22.04:7.2.3-complete
+    env:
+      PYTORCH_VERSION: "2.11.0"
+      ROCM_INDEX: "rocm7.2"
+    steps:
+    - uses: actions/checkout@v4
+    - name: Install Python and torch+rocm
+      run: |-
+        apt-get update
+        apt-get install -y --no-install-recommends python3 python3-dev python3-pip git
+        python3 -m pip install --upgrade pip
+        python3 -m pip install --index-url https://download.pytorch.org/whl/${ROCM_INDEX} torch==${PYTORCH_VERSION}
+    - name: Verify torch is ROCm-built
+      run: |-
+        python3 -c "import torch; assert torch.version.hip is not None, 'torch is not HIP-built'; print('torch.version.hip:', torch.version.hip)"
+    - name: Build pytorch3d _C extension (build only, no tests)
+      env:
+        # CPU-only runner: torch.cuda.is_available() is False, so force the
+        # CUDAExtension path. ROCM_HOME is auto-detected from /opt/rocm in
+        # the rocm/dev-ubuntu container.
+        FORCE_CUDA: "1"
+      run: |-
+        python3 -m pip install --no-build-isolation -v .
+    - name: Smoke import
+      # cd out of the checkout root so the source-tree pytorch3d/ directory
+      # (which has no _C.so since the build doesn't install in-place) doesn't
+      # shadow the site-packages install via sys.path[0] for `python -c`.
+      run: |-
+        cd /tmp
+        python3 -c "import torch; from pytorch3d import _C; print('PulsarRenderer:', hasattr(_C, 'PulsarRenderer')); print('n_symbols:', len([s for s in dir(_C) if not s.startswith('_')]))"

.gitignore

@@ -6,6 +6,17 @@ dist/
 **/.ipynb_checkpoints
 **/.ipynb_checkpoints/**
 
+# Build artifacts produced in-place by torch.utils.cpp_extension auto-hipify
+# when pytorch3d is built against a ROCm PyTorch.
+*.so
+pytorch3d/csrc/**/*.hip
+pytorch3d/csrc/**/*_hip.cpp
+pytorch3d/csrc/**/*_hip.h
+pytorch3d/csrc/**/*_hip.cuh
+
+# Debug PNG dumps written by pulsar tests when FB_TEST is unset.
+tests/pulsar/test_out/
+
 
 # Docusaurus site
 website/yarn.lock

pytorch3d/csrc/ext.cpp

@@ -98,7 +98,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("marching_cubes", &MarchingCubes);
 
   // Pulsar.
-  // Pulsar not enabled on AMD.
 #ifdef PULSAR_LOGGING_ENABLED
   c10::ShowLogInfoToStderr();
 #endif

pytorch3d/csrc/pulsar/gpu/commands.h

@@ -125,29 +125,45 @@ INLINE DEVICE float3 WARP_SUM_FLOAT3(
 // Floating point.
 // #define FMUL(a, b) __fmul_rn((a), (b))
 #define FMUL(a, b) ((a) * (b))
-#define FDIV(a, b) __fdiv_rn((a), (b))
 // #define FSUB(a, b) __fsub_rn((a), (b))
 #define FSUB(a, b) ((a) - (b))
+// HIP has no per-instruction rounding-mode override (round-to-nearest-even is
+// the default at the hardware level), so the CUDA *_rn intrinsics have no
+// runtime equivalent and we use plain operators. The HIP/clang compiler may
+// fuse a+b*c into a single-rounding FMA where CUDA's _rn would have prevented
+// it; if that becomes a numerical issue, add `-ffp-contract=off` to the pulsar
+// nvcc_args in setup.py.
+#if defined(USE_ROCM)
+#define FDIV(a, b) ((a) / (b))
+#define FADD(a, b) ((a) + (b))
+#define FSQRT(a) sqrtf(a)
+#define FPOW(a, b) powf((a), (b))
+#define FSATURATE(x) fmaxf(0.0f, fminf(1.0f, (x)))
+#define FMA(x, y, z) fmaf((x), (y), (z))
+#define FRCP(x) (1.0f / (x))
+#else
+#define FDIV(a, b) __fdiv_rn((a), (b))
 #define FADD(a, b) __fadd_rn((a), (b))
 #define FSQRT(a) __fsqrt_rn(a)
+#define FPOW(a, b) __powf((a), (b))
+#define FSATURATE(x) __saturatef(x)
+/** Calculates x*y+z. */
+#define FMA(x, y, z) __fmaf_rn((x), (y), (z))
+#define FRCP(x) __frcp_rn(x)
+#endif
 #define FEXP(a) fasterexp(a)
 #define FLN(a) fasterlog(a)
-#define FPOW(a, b) __powf((a), (b))
 #define FMAX(a, b) fmax((a), (b))
 #define FMIN(a, b) fmin((a), (b))
 #define FCEIL(a) ceilf(a)
 #define FFLOOR(a) floorf(a)
 #define FROUND(x) nearbyintf(x)
-#define FSATURATE(x) __saturatef(x)
 #define FABS(a) abs(a)
 #define IASF(a, loc) (loc) = __int_as_float(a)
 #define FASI(a, loc) (loc) = __float_as_int(a)
 #define FABSLEQAS(a, b, c) \
   ((a) <= (b) ? FSUB((b), (a)) <= (c) : FSUB((a), (b)) < (c))
-/** Calculates x*y+z. */
-#define FMA(x, y, z) __fmaf_rn((x), (y), (z))
 #define I2F(a) __int2float_rn(a)
-#define FRCP(x) __frcp_rn(x)
 #if !defined(USE_ROCM)
 __device__ static float atomicMax(float* address, float val) {
   int* address_as_i = (int*)address;
@@ -201,8 +217,16 @@ __device__ static float atomicMin(float* address, float val) {
   ATOMICADD(&((PTR)->x), VAL.x); \
   ATOMICADD(&((PTR)->y), VAL.y); \
   ATOMICADD(&((PTR)->z), VAL.z);
-#if (CUDART_VERSION >= 10000) && (__CUDA_ARCH__ >= 600)
+#if !defined(USE_ROCM) && (CUDART_VERSION >= 10000) && (__CUDA_ARCH__ >= 600)
 #define ATOMICADD_B(PTR, VAL) atomicAdd_block((PTR), (VAL))
+#elif defined(USE_ROCM)
+// HIP has no atomicAdd_block, but the semantic equivalent is a relaxed
+// fetch_add scoped to the workgroup (HIP's name for a CUDA thread block).
+// This avoids device-wide L2-coherent atomics for what are block-local
+// counters in pulsar's inner sphere-loading loop.
+#define ATOMICADD_B(PTR, VAL) \
+  __hip_atomic_fetch_add(     \
+      (PTR), (VAL), __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_WORKGROUP)
 #else
 #define ATOMICADD_B(PTR, VAL) ATOMICADD(PTR, VAL)
 #endif

pytorch3d/csrc/utils/warp_reduce.cuh

@@ -10,10 +10,20 @@
 #include <math.h>
 #include <cstdio>
 
-// Helper functions WarpReduceMin and WarpReduceMax used in .cu files
-// Starting in Volta, instructions are no longer synchronous within a warp.
-// We need to call __syncwarp() to sync the 32 threads in the warp
-// instead of all the threads in the block.
+// Helper functions WarpReduceMin and WarpReduceMax used in .cu files.
+// Starting in Volta, instructions are no longer synchronous within a warp,
+// so on CUDA __syncwarp() is required between dependent shared-memory
+// accesses in the unrolled tail reduction.
+//
+// On AMD/HIP no __syncwarp() is needed here: all wavefront lanes execute
+// in lockstep (AMD has no equivalent of NVIDIA's Independent Thread
+// Scheduling), and the AMDGPU memory model guarantees that LDS operations
+// issued by the same wavefront are observed in program order without an
+// explicit s_waitcnt — see the LLVM AMDGPU backend memory-model rules
+// (https://llvm.org/docs/AMDGPUUsage.html) and the HIP hardware-
+// implementation docs. This holds for both wave32 (RDNA, gfx10xx/11xx/
+// 12xx) and wave64 (CDNA, gfx9xx), so the USE_ROCM skip is
+// architecture-independent.
 
 template <typename scalar_t>
 __device__ void
@@ -23,8 +33,6 @@ WarpReduceMin(scalar_t* min_dists, int64_t* min_idxs, const size_t tid) {
     min_idxs[tid] = min_idxs[tid + 32];
     min_dists[tid] = min_dists[tid + 32];
   }
-// AMD does not use explicit syncwarp and instead automatically inserts memory
-// fences during compilation.
 #if !defined(USE_ROCM)
   __syncwarp();
 #endif

setup.py

@@ -14,7 +14,7 @@
 
 import torch
 from setuptools import find_packages, setup
-from torch.utils.cpp_extension import CppExtension, CUDA_HOME, CUDAExtension
+from torch.utils.cpp_extension import CppExtension, CUDA_HOME, CUDAExtension, ROCM_HOME
 
 
 def get_existing_ccbin(nvcc_args: List[str]) -> Optional[str]:
@@ -53,20 +53,25 @@ def get_extensions():
     define_macros = []
     include_dirs = [extensions_dir]
 
+    # ROCm/HIP support. When PyTorch is built with HIP, the cpp_extension
+    # BuildExtension auto-hipifies .cu sources and swaps nvcc -> hipcc.
+    is_rocm = torch.version.hip is not None
+
     force_cuda = os.getenv("FORCE_CUDA", "0") == "1"
     force_no_cuda = os.getenv("PYTORCH3D_FORCE_NO_CUDA", "0") == "1"
+    gpu_home_available = CUDA_HOME is not None or (is_rocm and ROCM_HOME is not None)
     if (
-        not force_no_cuda and torch.cuda.is_available() and CUDA_HOME is not None
+        not force_no_cuda and torch.cuda.is_available() and gpu_home_available
     ) or force_cuda:
         extension = CUDAExtension
+
         sources += source_cuda
         define_macros += [("WITH_CUDA", None)]
         # Thrust is only used for its tuple objects.
         # With CUDA 11.0 we can't use the cudatoolkit's version of cub.
         # We take the risk that CUB and Thrust are incompatible, because
         # we aren't using parts of Thrust which actually use CUB.
         define_macros += [("THRUST_IGNORE_CUB_VERSION_CHECK", None)]
-        cub_home = os.environ.get("CUB_HOME", None)
         nvcc_args = [
             "-DCUDA_HAS_FP16=1",
             "-D__CUDA_NO_HALF_OPERATORS__",
@@ -76,35 +81,40 @@ def get_extensions():
         if os.name != "nt":
             nvcc_args.append("-std=c++17")
 
-        # CUDA 13.0+ compatibility flags for pulsar.
-        # Starting with CUDA 13, __global__ function visibility changed.
-        # See: https://developer.nvidia.com/blog/
-        #      cuda-c-compiler-updates-impacting-elf-visibility-and-linkage/
-        cuda_version = torch.version.cuda
-        if cuda_version is not None:
-            major = int(cuda_version.split(".")[0])
-            if major >= 13:
-                nvcc_args.extend(
-                    [
-                        "--device-entity-has-hidden-visibility=false",
-                        "-static-global-template-stub=false",
-                    ]
+        if not is_rocm:
+            # CUDA 13.0+ compatibility flags for pulsar.
+            # Starting with CUDA 13, __global__ function visibility changed.
+            # See: https://developer.nvidia.com/blog/
+            #      cuda-c-compiler-updates-impacting-elf-visibility-and-linkage/
+            cuda_version = torch.version.cuda
+            if cuda_version is not None:
+                major = int(cuda_version.split(".")[0])
+                if major >= 13:
+                    nvcc_args.extend(
+                        [
+                            "--device-entity-has-hidden-visibility=false",
+                            "-static-global-template-stub=false",
+                        ]
+                    )
+
+            # NVIDIA CUB. On ROCm, hipcub from the ROCm toolchain is used and
+            # no external CUB_HOME is required.
+            cub_home = os.environ.get("CUB_HOME", None)
+            if cub_home is None:
+                prefix = os.environ.get("CONDA_PREFIX", None)
+                if prefix is not None and os.path.isdir(prefix + "/include/cub"):
+                    cub_home = prefix + "/include"
+
+            if cub_home is None:
+                warnings.warn(
+                    "The environment variable `CUB_HOME` was not found. "
+                    "NVIDIA CUB is required for compilation and can be downloaded "
+                    "from `https://github.com/NVIDIA/cub/releases`. You can unpack "
+                    "it to a location of your choice and set the environment variable "
+                    "`CUB_HOME` to the folder containing the `CMakeListst.txt` file."
                 )
-        if cub_home is None:
-            prefix = os.environ.get("CONDA_PREFIX", None)
-            if prefix is not None and os.path.isdir(prefix + "/include/cub"):
-                cub_home = prefix + "/include"
-
-        if cub_home is None:
-            warnings.warn(
-                "The environment variable `CUB_HOME` was not found. "
-                "NVIDIA CUB is required for compilation and can be downloaded "
-                "from `https://github.com/NVIDIA/cub/releases`. You can unpack "
-                "it to a location of your choice and set the environment variable "
-                "`CUB_HOME` to the folder containing the `CMakeListst.txt` file."
-            )
-        else:
-            include_dirs.append(os.path.realpath(cub_home).replace("\\ ", " "))
+            else:
+                include_dirs.append(os.path.realpath(cub_home).replace("\\ ", " "))
         nvcc_flags_env = os.getenv("NVCC_FLAGS", "")
         if nvcc_flags_env != "":
             nvcc_args.extend(nvcc_flags_env.split(" "))
@@ -113,7 +123,9 @@ def get_extensions():
         # https://github.com/facebookresearch/pytorch3d/issues/436
         # It is harmless after https://github.com/pytorch/pytorch/pull/47404 .
         # But it can be problematic in torch 1.7.0 and 1.7.1
-        if torch.__version__[:4] != "1.7.":
+        # On ROCm the host compiler is selected by hipcc itself; -ccbin is
+        # an nvcc-only flag.
+        if not is_rocm and torch.__version__[:4] != "1.7.":
             CC = os.environ.get("CC", None)
             if CC is not None:
                 existing_CC = get_existing_ccbin(nvcc_args)

tests/test_cameras_alignment.py

@@ -108,9 +108,14 @@ def _corresponding_cameras_alignment_test_case(
             cameras, cameras_tgt, estimate_scale=estimate_scale, mode=mode
         )
 
-        if batch_size <= 2 and mode == "centers":
-            # underdetermined case - check only the center alignment error
-            # since the rotation and translation are ambiguous here
+        if batch_size <= 3 and mode == "centers":
+            # Underdetermined case: with <= 3 camera centers in 3D, the points
+            # span at most a 2D subspace after mean-centering, so the Umeyama
+            # SVD has a zero (or near-zero) third singular value and the
+            # rotation around the degenerate axis is ambiguous. Different
+            # SVD implementations (e.g. rocBLAS on RDNA vs CDNA, or
+            # cuBLAS) make different valid choices in that null direction.
+            # Only the camera centers are well-defined here, so check those.
             self.assertClose(
                 cameras_aligned.get_camera_center(),
                 cameras_tgt.get_camera_center(),

tests/test_point_mesh_distance.py

@@ -707,9 +707,14 @@ def test_point_face_distance(self):
 
         # Compare
         self.assertClose(grad_points_naive.cpu(), grad_points_cuda.cpu(), atol=1e-7)
-        self.assertClose(grad_faces_naive, grad_faces_cuda.cpu(), atol=5e-7)
+        # DistanceBackward uses atomicAdd to accumulate gradients into the
+        # face buffer and explicitly calls alertNotDeterministic; the FP add
+        # order differs across GPU architectures (e.g. between 32- and
+        # 64-lane warps), producing tiny rounding differences. Use the same
+        # 5e-6 tolerance as test_face_point_distance below.
+        self.assertClose(grad_faces_naive, grad_faces_cuda.cpu(), atol=5e-6)
         self.assertClose(grad_points_naive.cpu(), grad_points_cpu, atol=1e-7)
-        self.assertClose(grad_faces_naive, grad_faces_cpu, atol=5e-7)
+        self.assertClose(grad_faces_naive, grad_faces_cpu, atol=5e-6)
 
     def test_face_point_distance(self):
         """

tests/test_points_alignment.py

@@ -669,12 +669,20 @@ def align_and_get_mse(weights_):
                     desired_det *= -1.0
                 self._assert_all_close(torch.det(R_est), desired_det, msg, w, atol=2e-5)
 
-            # check that the transformed point cloud
-            # X matches X_t
-            X_t_est = _apply_pcl_transformation(X, R_est, T_est, s=s_est)
-            self._assert_all_close(
-                X_t, X_t_est, assert_error_message, w[:, None, None], atol=2e-5
-            )
+                # check that the transformed point cloud
+                # X matches X_t.
+                # Only valid when the problem setup is unambiguous: when
+                # n_points <= dim the centered point cloud is rank-deficient
+                # and the rotation around the degenerate axis is determined
+                # only by the SVD's null-space convention, which differs
+                # across BLAS implementations (e.g. rocBLAS on RDNA vs CDNA,
+                # or cuBLAS). Applying any of those valid rotations to the
+                # uncentered X yields a different X_t_est even though the
+                # algorithm is correct.
+                X_t_est = _apply_pcl_transformation(X, R_est, T_est, s=s_est)
+                self._assert_all_close(
+                    X_t, X_t_est, assert_error_message, w[:, None, None], atol=2e-5
+                )
 
     def _assert_all_close(self, a_, b_, err_message, weights=None, atol=1e-6):
         if isinstance(a_, Pointclouds):