Fix scalar rasterizer packet iteration (#1569)

Summary:
Pull Request resolved: #1569

PyMomentum renderer tests were skipped in OSS CI because the rasterizer crashed
when built without SIMD. The failing path was `pymomentum.renderer` calling
software rasterization helpers such as `rasterize_spheres`, which eventually
iterate rasterizer primitives with `drjit::range<IntP>()`.

This works in internal Buck builds because those builds use SIMD compiler flags,
so Dr.Jit sees `IntP` as a packed SIMD packet and `drjit::range<IntP>()` produces
packet indices as intended. OSS/PyPI-style builds can compile with SIMD disabled,
where `IntP` becomes scalar (`Packet<int, 1>`). In that mode `drjit::range<IntP>()`
goes through Dr.Jit's recursive static-array range implementation instead of the
packed SIMD range path, and the one-dimensional scalar case underflows internally.
That produces invalid indices and can crash the renderer.

Fix this by adding a rasterizer-local `intPacketRange()` helper:

- SIMD builds continue to use `drjit::range<IntP>()`, with a static assertion that
  `IntP` is actually packed.
- Scalar/non-SIMD builds use explicit packet iteration based on `drjit::arange<IntP>()`
  and a tail mask.

Update all rasterizer primitive loops to use the helper, then unskip the affected
Linux and macOS PyMomentum renderer tests in OSS CI. The Windows `test_lighting`
skip remains because it tracks a separate CMake stack-overflow issue.

Reviewed By: cdtwigg

Differential Revision: D111828395

fbshipit-source-id: 556184370c9decba521cde0e924b6a0e69330a0d

Author: yutingye

CMakeLists.txt

@@ -261,6 +261,13 @@ if(MOMENTUM_BUILD_TESTING OR MOMENTUM_BUILD_PYMOMENTUM)
   set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 endif()
 
+if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64|ARM64)$")
+  # GCC emits very noisy psABI notes for arm64 vector and Eigen types after its
+  # GCC 10 ABI correction. We build the wheel with one toolchain, so these notes
+  # do not indicate a mixed-ABI problem and only obscure real warnings in CI.
+  add_compile_options($<$<COMPILE_LANGUAGE:CXX>:-Wno-psabi>)
+endif()
+
 if(MSVC)
   add_compile_options(/bigobj)
   add_compile_options(/utf-8)

momentum/rasterizer/image.cpp

@@ -166,7 +166,7 @@ void alphaMatte(Span2f zBuffer, Span3f rgbBuffer, Span<T, 3> tgtImage, float alp
       drjit::scatter(tempBuffer.data(), rgba, colIndices);
     }
 
-    for (const auto [colIndices, colMask] : drjit::range<IntP>(imageWidthDownsampled)) {
+    for (const auto [colIndices, colMask] : intPacketRange(imageWidthDownsampled)) {
       auto rgba = drjit::zeros<Vector4fP>();
       // Walk through the columns and sum horizontally
       for (int32_t offset = 0; offset < downsampleAmount; ++offset) {

momentum/rasterizer/rasterizer.cpp

@@ -321,7 +321,7 @@ void rasterizeLinesImp(
 
   const Vector3f color_drjit = toEnokiVec(color);
 
-  for (auto [lineIndices, lineMask] : drjit::range<IntP>(nLines)) {
+  for (auto [lineIndices, lineMask] : intPacketRange(nLines)) {
     const auto lineStart_world =
         drjit::gather<Vector3fP>(positions_world.data(), 2 * lineIndices + 0, lineMask);
     const auto lineEnd_world =
@@ -407,7 +407,7 @@ void rasterizeCirclesImp(
   const int combinedRadiusPixels =
       std::clamp<int>(static_cast<int>(std::ceil(combinedRadius)), 1, cameraSimd.imageHeight() - 1);
 
-  for (auto [circleIndices, circleMask] : drjit::range<IntP>(nCircles)) {
+  for (auto [circleIndices, circleMask] : intPacketRange(nCircles)) {
     const auto center_world =
         drjit::gather<Vector3fP>(positions_world.data(), circleIndices, circleMask);
     const auto [center_window, validProj] = cameraSimd.worldToWindow(center_world);
@@ -490,7 +490,7 @@ void rasterizeWireframeImp(
 
   const Vector3f color_drjit = toEnokiVec(color);
 
-  for (auto [triangleIndices, triangleMask] : drjit::range<IntP>(nTriangles)) {
+  for (auto [triangleIndices, triangleMask] : intPacketRange(nTriangles)) {
     auto triangles_cur = drjit::gather<Vector3iP>(triangles.data(), triangleIndices, triangleMask);
     std::array<Vector3fP, 3> p_tri_window;
     IntP::MaskType validTriangles = triangleMask;

momentum/rasterizer/rasterizer_triangles.cpp

@@ -562,7 +562,7 @@ void rasterizeMeshImp(
   int* vertexIndexBufferPtr = dataOrNull(vertexIndexBuffer);
   int* triangleIndexBufferPtr = dataOrNull(triangleIndexBuffer);
 
-  for (auto [triangleIndices, triangleMask] : drjit::range<IntP>(nTriangles)) {
+  for (auto [triangleIndices, triangleMask] : intPacketRange(nTriangles)) {
     auto triangles_cur = drjit::gather<Vector3iP>(triangles.data(), triangleIndices, triangleMask);
     std::array<Vector3fP, 3> p_tri_window;
     Matrix3fP p_tri_eye;
@@ -746,7 +746,7 @@ void rasterizeSplatsImp(
   const std::array<Vector3f, 4> quadTextureCoords = {
       Vector3f(-1, -1, 0), Vector3f(1, -1, 0), Vector3f(1, 1, 0), Vector3f(-1, 1, 0)};
 
-  for (auto [splatIndices, splatMask] : drjit::range<IntP>(nSplats)) {
+  for (auto [splatIndices, splatMask] : intPacketRange(nSplats)) {
     auto position_world = drjit::gather<Vector3fP>(positions_world.data(), splatIndices, splatMask);
     auto normal_world = drjit::gather<Vector3fP>(normals_world.data(), splatIndices, splatMask);
 

momentum/rasterizer/utility.h

@@ -10,12 +10,14 @@
 #include <drjit/array.h>
 #include <drjit/fwd.h>
 #include <drjit/matrix.h>
+#include <drjit/util.h>
 #include <mdspan/mdspan.hpp>
 #include <momentum/common/exception.h>
 #include <momentum/rasterizer/fwd.h>
 #include <Eigen/Core>
 #include <Eigen/Geometry>
 #include <array>
+#include <limits>
 #include <memory>
 #include <sstream>
 #include <utility>
@@ -25,7 +27,76 @@ namespace momentum::rasterizer {
 
 using index_t = std::ptrdiff_t;
 
-// Actual implementations
+// drjit::range<IntP>() assumes IntP is a packed SIMD type. Use it only in
+// SIMD builds; scalar/non-SIMD builds need explicit packet iteration.
+#if defined(MOMENTUM_ENABLE_SIMD) && MOMENTUM_ENABLE_SIMD
+/// Returns integer packet indices and validity masks for rasterizer primitive loops.
+inline auto intPacketRange(index_t count) {
+  static_assert(IntP::IsPacked, "MOMENTUM_ENABLE_SIMD requires packed Dr.Jit IntP packets.");
+  return drjit::range<IntP>(count);
+}
+#else
+/// Scalar fallback range yielding IntP index packets and validity masks.
+class ScalarIntPacketRange {
+ public:
+  /// Creates a range over integer indices in [0, count).
+  explicit ScalarIntPacketRange(index_t count) : count_(count) {
+    MT_THROW_IF(
+        count < 0 || count > std::numeric_limits<int>::max(),
+        "Packet range count {} is outside the supported int range.",
+        count);
+  }
+
+  /// Iterator for ScalarIntPacketRange.
+  class Iterator {
+   public:
+    /// Creates an iterator at a packet base offset.
+    Iterator(index_t base, index_t count) : base_(base), count_(count) {}
+
+    /// Returns true when the iterator has not reached the end sentinel.
+    bool operator!=(const Iterator& other) const {
+      return base_ != other.base_;
+    }
+
+    /// Advances to the next packet of indices.
+    Iterator& operator++() {
+      base_ += static_cast<index_t>(kSimdPacketSize);
+      return *this;
+    }
+
+    /// Returns the current index packet and a mask for lanes below count.
+    auto operator*() const {
+      const IntP indices = IntP(static_cast<int>(base_)) + drjit::arange<IntP>();
+      const auto mask = indices < IntP(static_cast<int>(count_));
+      return std::pair{indices, mask};
+    }
+
+   private:
+    index_t base_;
+    index_t count_;
+  };
+
+  /// Returns an iterator for the first packet.
+  Iterator begin() const {
+    return Iterator(0, count_);
+  }
+
+  /// Returns the sentinel iterator after the final packet.
+  Iterator end() const {
+    constexpr index_t kPacketSize = static_cast<index_t>(kSimdPacketSize);
+    const index_t endBase = ((count_ + kPacketSize - 1) / kPacketSize) * kPacketSize;
+    return Iterator(endBase, count_);
+  }
+
+ private:
+  index_t count_;
+};
+
+/// Returns integer packet indices and validity masks for scalar rasterizer builds.
+inline ScalarIntPacketRange intPacketRange(index_t count) {
+  return ScalarIntPacketRange(count);
+}
+#endif
 
 template <std::size_t R>
 std::string formatTensorSizes(const std::array<index_t, R>& extents) {

pixi.toml

@@ -328,8 +328,7 @@ build_py = { cmd = "pip install . -v --no-build-isolation", env = { FBXSDK_PATH
 """, MOMENTUM_ENABLE_SIMD = "ON", TORCH_CUDA_ARCH_LIST = "5.0;6.0;6.1;7.0;7.5;8.0;8.6;8.9;9.0+PTX" }, depends-on = [
     "install_deps",
 ] }
-# TODO(T219163027): Skip TestRendering entirely on Linux — drjit segfaults in scalar SIMD mode (test_lighting, test_rasterize_*, etc.)
-test_py = { cmd = "pytest pymomentum/test/ -k 'not TestRendering'", env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
+test_py = { cmd = "pytest pymomentum/test/", env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
     "build_py",
 ] }
 doc_py = { cmd = "sphinx-build -a -E -W -b html pymomentum/doc build/python_api_doc", depends-on = [
@@ -356,8 +355,7 @@ build_py = { cmd = "pip install . -v --no-build-isolation", env = { CMAKE_ARGS =
     -DBUILD_SHARED_LIBS=OFF \
     -DCMAKE_CXX_FLAGS=-fsigned-char
 """, MOMENTUM_ENABLE_SIMD = "ON" } }
-# TODO(T219163027): Skip TestRendering entirely on Linux — drjit segfaults in scalar SIMD mode
-test_py = { cmd = "pytest pymomentum/test/ -k 'not TestRendering'", env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
+test_py = { cmd = "pytest pymomentum/test/", env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
     "build_py",
 ] }
 
@@ -380,8 +378,7 @@ build_py = { cmd = "pip install . -v", env = { CMAKE_ARGS = """
     -DBUILD_SHARED_LIBS=OFF
 """, MOMENTUM_ENABLE_FBX_SAVING = "OFF", MOMENTUM_ENABLE_SIMD = "ON" } }
 # TODO: Remove -k option, once saving tests are disabled programmatically when Autodesk FBX SDK is unavailable (FBX loading always works via OpenFBX)
-# TODO(T219163027): Remove test_lighting skip once drjit::range<IntP> crash in scalar SIMD mode is fixed
-test_py = { cmd = "pytest pymomentum/test/ -k 'not ((TestFBXIO and (test_save_motions_with_joint_params or test_save_motions_with_model_params)) or test_lighting)'", env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
+test_py = { cmd = "pytest pymomentum/test/ -k 'not (TestFBXIO and (test_save_motions_with_joint_params or test_save_motions_with_model_params))'", env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
     "build_py", "install",
 ] }
 doc_py = { cmd = "sphinx-build -a -E -W -b html pymomentum/doc build/python_api_doc", depends-on = [
@@ -411,8 +408,7 @@ build_py = { cmd = "pip install . -v", env = { CMAKE_ARGS = """
     -DBUILD_SHARED_LIBS=OFF
 """, MOMENTUM_ENABLE_FBX_SAVING = "OFF", MOMENTUM_ENABLE_SIMD = "ON" } }
 # TODO: Remove -k option, once saving tests are disabled programmatically when Autodesk FBX SDK is unavailable (FBX loading always works via OpenFBX)
-# TODO(T219163027): Remove test_lighting skip once drjit::range<IntP> crash in scalar SIMD mode is fixed
-test_py = { cmd = "pytest pymomentum/test/ -k 'not ((TestFBXIO and (test_save_motions_with_joint_params or test_save_motions_with_model_params)) or test_lighting)'", env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
+test_py = { cmd = "pytest pymomentum/test/ -k 'not (TestFBXIO and (test_save_motions_with_joint_params or test_save_motions_with_model_params))'", env = { MOMENTUM_MODELS_PATH = "momentum/" }, depends-on = [
     "build_py", "install",
 ] }
 doc_py = { cmd = "sphinx-build -a -E -W -b html pymomentum/doc build/python_api_doc", depends-on = [