Port pulsar to ROCm/HIP

Re-enables the pulsar subrenderer on ROCm-built PyTorch. Pulsar was excluded from the ROCm build in the parent commit (PR #2039) because its 74 source files use a mix of CUDA intrinsics with no direct HIP equivalent. This commit makes it build and run on AMD MI250X (gfx90a, warpSize=64, HIP 7.2) using only HIP runtime APIs and standard math — no `amdgcn` inline assembly.

Most of the lift is build-system reversal. The pulsar source filter in `setup.py`, the `#ifndef USE_ROCM` wraps around the pulsar `#include`s and pybind block in `pytorch3d/csrc/ext.cpp`, and the `hasattr(_C, "PulsarRenderer")` guards in the three `pytorch3d/renderer/.../__init__.py` files all go away. The only C++ surgery is in `pytorch3d/csrc/pulsar/gpu/commands.h`:

The CUDA `_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 — the rounding mode lives in the wavefront `MODE` register, defaulting to round-to-nearest-even. So under `USE_ROCM` these macros expand to plain operators / `sqrtf` / `fmaf` / `1.0f/x` / `fmaxf(0,fminf(1,x))`, which are rounding-mode-equivalent on HIP. The HIP/clang compiler may fuse `a+b*c` into a single-rounding FMA where CUDA's `_rn` would have prevented it; the upstream pulsar authors already accept compiler-discretion FMA fusion for `FMUL`/`FSUB` (the `_rn` variants are commented out in the source), and pulsar's tests pass under the same trade-off on HIP. If FMA-fusion drift ever surfaces as a numerical issue, add `-ffp-contract=off` to pulsar's HIPCC flags.

`__powf` is replaced with `powf` for clarity on the ROCm arm.

`atomicAdd_block` has no HIP function-name equivalent, but the *semantic* equivalent is `__hip_atomic_fetch_add(ptr, val, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_WORKGROUP)`. Plain HIP `atomicAdd` is device-scope (`__HIP_MEMORY_SCOPE_AGENT`, see `amd_hip_atomic.h:218`), which is strictly stronger than block-scope — correct but forces L2-coherent atomics on what are block-local counters at `renderer.render.device.h:186,297` (sphere-loading counter and per-pixel-done counter). The workgroup-scoped builtin gets us back to the cheaper LDS atomic path that CUDA's `atomicAdd_block` uses.

Verified absent-from-HIP and explicitly replaced: `__fdiv_rn`, `__fadd_rn`, `__fsqrt_rn`, `__fmaf_rn`, `__frcp_rn`, `__saturatef`, `atomicAdd_block` (semantic replacement via `__hip_atomic_fetch_add` + `__HIP_MEMORY_SCOPE_WORKGROUP`).

Verified present-in-HIP and left alone: `hsqrt` (`amd_hip_fp16.h`), `__int2float_rn` (`amd_device_functions.h:536`, defined as `(float)x`), `__hadd`/`__hsub2`/`__hmul2` (`amd_hip_fp16.h`), `__clz`/`__popc`/`__popcll`/`__mul24` (`amd_device_functions.h`), native `atomicMin(float*,float)`/`atomicMax(float*,float)` (`amd_hip_atomic.h`), and the existing `__HIP_PLATFORM_AMD__` `__ballot`+`__popcll` warpSize=64 workaround at `include/renderer.render.device.h:289-295`.

The unreferenced 32-lane WARP_CUMSUM/WARP_MAX/WARP_SUM helpers at `gpu/commands.h:66-123` remain gated out on ROCm — the TODO comment about ROCM-6.2 was technically right that HIP supports `__shfl_*_sync` now, but since nothing in pulsar actually calls these helpers, there's no live warpSize=32 assumption to fix.

Test results on AMD Instinct MI250X (gfx90a, HIP 7.2):

  tests/pulsar/test_forward.py                                          5 passed
  tests/pulsar/test_channels.py                                         1 passed
  tests/pulsar/test_depth.py                                            1 passed
  tests/pulsar/test_hands.py                                            1 passed
  tests/pulsar/test_ortho.py                                            1 passed
  tests/pulsar/test_small_spheres.py                                    1 passed
  tests/test_render_points.py::test_simple_sphere_pulsar                1 passed
  tests/test_render_points.py::test_unified_inputs_pulsar               1 passed
  tests/test_camera_conversions.py::test_pulsar_conversion              1 passed

Pulsar tests need `FB_TEST=1` because `tests/pulsar/test_forward.py::test_principal_point` writes a 1-channel debug PNG via `imageio.imsave`, and current `imageio`'s pillow backend refuses 1-channel writes (`ValueError: Can't write images with one color channel.`). The test gates the imsave block on `not os.environ.get("FB_TEST", False)`; the actual GPU render and `np.allclose` assertion that follow have nothing to do with imageio. This is an imageio/pillow versioning issue, not the HIP port.

Parent PR's 123 core CUDA-kernel tests still pass on the same host — no regression.

.gitignore now also excludes `tests/pulsar/test_out/` (PNG debug dumps written when FB_TEST is unset).

Co-Authored-By: Claude Opus 4.7 (1M context)

Author: jeffdaily

.gitignore

@@ -14,6 +14,9 @@ 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

@@ -6,15 +6,11 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// Pulsar is not currently portable to ROCm/HIP and is excluded from the
-// build there. Tracked as a separate follow-up port.
-#ifndef USE_ROCM
 // clang-format off
 #include "./pulsar/global.h" // Include before <torch/extension.h>.
 // clang-format on
 #include "./pulsar/pytorch/renderer.h"
 #include "./pulsar/pytorch/tensor_util.h"
-#endif // !USE_ROCM
 #include "ball_query/ball_query.h"
 #include "blending/sigmoid_alpha_blend.h"
 #include "compositing/alpha_composite.h"
@@ -102,8 +98,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("marching_cubes", &MarchingCubes);
 
   // Pulsar.
-  // Pulsar not enabled on AMD.
-#ifndef USE_ROCM
 #ifdef PULSAR_LOGGING_ENABLED
   c10::ShowLogInfoToStderr();
 #endif
@@ -189,5 +183,4 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.attr("MAX_UINT") = py::int_(MAX_UINT);
   m.attr("MAX_USHORT") = py::int_(MAX_USHORT);
   m.attr("PULSAR_MAX_GRAD_SPHERES") = py::int_(MAX_GRAD_SPHERES);
-#endif // !USE_ROCM
 }

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,17 @@ __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/renderer/__init__.py

@@ -75,16 +75,9 @@
     PointsRasterizationSettings,
     PointsRasterizer,
     PointsRenderer,
+    PulsarPointsRenderer,
     rasterize_points,
 )
-
-# PulsarPointsRenderer is only exposed when the native pulsar renderer was
-# built into pytorch3d._C (CUDA builds). On ROCm builds, pulsar is excluded
-# and the symbol simply will not exist on this module.
-try:
-    from .points import PulsarPointsRenderer  # noqa: F401
-except ImportError:
-    pass
 from .splatter_blend import SplatterBlender
 from .utils import (
     convert_to_tensors_and_broadcast,

pytorch3d/renderer/points/__init__.py

@@ -8,19 +8,11 @@
 
 import torch
 
-from pytorch3d import _C
-
 from .compositor import AlphaCompositor, NormWeightedCompositor
+from .pulsar.unified import PulsarPointsRenderer
 from .rasterize_points import rasterize_points
 from .rasterizer import PointsRasterizationSettings, PointsRasterizer
 from .renderer import PointsRenderer
 
 
-# PulsarPointsRenderer wraps the native pulsar renderer, which is not yet
-# available on ROCm builds. Only expose the wrapper when the native class is
-# present so that ``import pytorch3d.renderer`` succeeds on AMD GPUs.
-if hasattr(_C, "PulsarRenderer"):
-    from .pulsar.unified import PulsarPointsRenderer  # noqa: F401
-
-
 __all__ = [k for k in globals().keys() if not k.startswith("_")]

pytorch3d/renderer/points/pulsar/__init__.py

@@ -6,13 +6,4 @@
 
 # pyre-unsafe
 
-from pytorch3d import _C
-
-
-# The pulsar native renderer is not yet ported to ROCm/HIP; the C++ extension
-# omits PulsarRenderer when built for AMD. Importing the Python wrapper would
-# crash at class-definition time on those builds (it references _C.MAX_UINT /
-# _C.PulsarRenderer at module load). Gate the wrapper symbol on availability
-# of the native class so the rest of pytorch3d.renderer remains importable.
-if hasattr(_C, "PulsarRenderer"):
-    from .renderer import Renderer  # noqa: F401
+from .renderer import Renderer  # noqa: F401

setup.py

@@ -59,9 +59,7 @@ def get_extensions():
     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. Pulsar
-    # is not currently portable to ROCm and is excluded from the build; see
-    # the comment in pytorch3d/csrc/ext.cpp ("Pulsar not enabled on AMD").
+    # 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"
@@ -74,13 +72,6 @@ def get_extensions():
     ) or force_cuda:
         extension = CUDAExtension
 
-        if is_rocm:
-            pulsar_prefix = os.path.join(extensions_dir, "pulsar") + os.sep
-            sources = [s for s in sources if not s.startswith(pulsar_prefix)]
-            source_cuda = [
-                s for s in source_cuda if not s.startswith(pulsar_prefix)
-            ]
-
         sources += source_cuda
         define_macros += [("WITH_CUDA", None)]
         # Thrust is only used for its tuple objects.