replace IASF/FASI pulsar macros with inline bit-cast functions

Summary: This avoids UB as reported in https://github.com/facebookresearch/pytorch3d/pull/2035

Reviewed By: MichaelRamamonjisoa

Differential Revision: D113009614

fbshipit-source-id: 92e3c6fe34731104ce85de909892b292b4464160

pytorch3d/common/workaround/symeig3x3.py

@@ -77,6 +77,20 @@ class _SymEig3x3(nn.Module):
         if inputs.shape[-2:] != (3, 3):
             raise ValueError("Only inputs of shape (..., 3, 3) are supported.")
 
+        # The thresholds and regularizers used below are absolute, so they are
+        # only meaningful for inputs of order 1: rescale each matrix to unit
+        # maximum magnitude and undo it on the eigenvalues. The clamp only
+        # matters for the all-zero matrix. The scale is detached because the
+        # decomposition is homogeneous of degree 1, so its Jacobian is
+        # homogeneous of degree 0 and holding the scale constant is exact.
+        scale = (
+            inputs.detach()
+            .abs()
+            .amax(dim=(-2, -1), keepdim=True)
+            .clamp(min=torch.finfo(inputs.dtype).tiny)
+        )
+        inputs = inputs / scale
+
         inputs_diag = inputs.diagonal(dim1=-2, dim2=-1)
         inputs_trace = inputs_diag.sum(-1)
         q = inputs_trace / 3.0
@@ -120,7 +134,7 @@ class _SymEig3x3(nn.Module):
         else:
             eigenvecs = None
 
-        return eigenvals, eigenvecs
+        return eigenvals * scale.squeeze(-1), eigenvecs
 
     def _construct_eigenvecs_set(
         self, inputs: torch.Tensor, eigenvals: torch.Tensor

pytorch3d/csrc/pulsar/gpu/commands.h

@@ -159,8 +159,6 @@ INLINE DEVICE float3 WARP_SUM_FLOAT3(
 #define FFLOOR(a) floorf(a)
 #define FROUND(x) nearbyintf(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))
 #define I2F(a) __int2float_rn(a)

pytorch3d/csrc/pulsar/host/commands.h

@@ -114,8 +114,6 @@ INLINE float atomicMin(float* address, float val) {
 #define FMA(x, y, z) ((x) * (y) + (z))
 #define I2F(a) static_cast<float>(a)
 #define FRCP(x) (1.f / (x))
-#define IASF(x, loc) memcpy(&(loc), &(x), sizeof(x))
-#define FASI(x, loc) memcpy(&(loc), &(x), sizeof(x))
 #define DMAX(a, b) std::max((a), (b))
 #define DMIN(a, b) std::min((a), (b))
 #define DSATURATE(a) DMIN(1., DMAX(0., (a)))

pytorch3d/csrc/pulsar/include/commands.h

@@ -9,7 +9,11 @@
 #ifndef PULSAR_NATIVE_COMMANDS_ROUTING_H_
 #define PULSAR_NATIVE_COMMANDS_ROUTING_H_
 
+#include <cstring>
 #include "../global.h"
+#if __cplusplus >= 202002L
+#include <bit>
+#endif
 
 // Commands available everywhere.
 #define MALLOC_HOST(VAR, TYPE, SIZE) \
@@ -32,4 +36,37 @@
 #include "../host/commands.h"
 #endif
 
+// Host + device int <-> float bit punning. Replaces IASF / FASI macros.
+IHD float pulsar_int_as_float(int v) noexcept {
+#ifdef __CUDA_ARCH__
+  return __int_as_float(v);
+#elif defined(__HIP_DEVICE_COMPILE__)
+  return __int_as_float(v);
+#else
+#if __cpp_lib_bit_cast >= 201806L
+  return std::bit_cast<float>(v);
+#else
+  float f;
+  std::memcpy(&f, &v, sizeof(f));
+  return f;
+#endif
+#endif
+}
+
+IHD int pulsar_float_as_int(float v) noexcept {
+#ifdef __CUDA_ARCH__
+  return __float_as_int(v);
+#elif defined(__HIP_DEVICE_COMPILE__)
+  return __float_as_int(v);
+#else
+#if __cpp_lib_bit_cast >= 201806L
+  return std::bit_cast<int>(v);
+#else
+  int i;
+  std::memcpy(&i, &v, sizeof(i));
+  return i;
+#endif
+#endif
+}
+
 #endif

pytorch3d/csrc/pulsar/include/renderer.calc_gradients.device.h

@@ -121,8 +121,8 @@ GLOBAL void calc_gradients(
       grad_im_l[2]);
   // Start processing.
   for (int grad_idx = 0; grad_idx < n_track; ++grad_idx) {
-    int sphere_idx;
-    FASI(forw_info_d[fwi_loc + 3 + 2 * grad_idx], sphere_idx);
+    int sphere_idx =
+        pulsar_float_as_int(forw_info_d[fwi_loc + 3 + 2 * grad_idx]);
     PASSERT(
         sphere_idx == -1 ||
         sphere_idx >= 0 && static_cast<uint>(sphere_idx) < num_balls);

pytorch3d/csrc/pulsar/include/renderer.fill_bg.device.h

@@ -41,7 +41,8 @@ GLOBAL void fill_bg(
     // sphere IDs and intersection depths.
     for (int i = 0; i < renderer.n_track; ++i) {
       int sphere_id = -1;
-      IASF(sphere_id, renderer.forw_info_d[write_loc + 3 + i * 2]);
+      renderer.forw_info_d[write_loc + 3 + i * 2] =
+          pulsar_int_as_float(sphere_id);
       renderer.forw_info_d[write_loc + 3 + i * 2 + 1] = -1.f;
     }
     if (mode == 0) {

pytorch3d/csrc/pulsar/include/renderer.render.device.h

@@ -406,7 +406,7 @@ GLOBAL void render(
         IMIN(n_track, tracker.get_n_hits()));
     for (int i = 0; i < n_track; ++i) {
       int sphere_id = tracker.get_closest_sphere_id(i);
-      IASF(sphere_id, forw_info_d[write_loc + 3 + i * 2]);
+      forw_info_d[write_loc + 3 + i * 2] = pulsar_int_as_float(sphere_id);
       forw_info_d[write_loc + 3 + i * 2 + 1] =
           tracker.get_closest_sphere_depth(i) == MAX_FLOAT
           ? -1.f

pytorch3d/io/obj_io.py

@@ -712,6 +712,7 @@ def save_obj(
         faces_normals_idx: LongTensor of shape (F, 3) giving the index into
             normals for each vertex in the face.
         verts_uvs: FloatTensor of shape (V, 2) giving the uv coordinate per vertex.
+            When provided with faces_uvs, UVs are saved even without a texture_map.
         faces_uvs: LongTensor of shape (F, 3) giving the index into verts_uvs for
             each vertex in the face.
         texture_map: FloatTensor of shape (H, W, 3) representing the texture map
@@ -757,7 +758,8 @@ def save_obj(
     if path_manager is None:
         path_manager = PathManager()
 
-    save_texture = all(t is not None for t in [faces_uvs, verts_uvs, texture_map])
+    save_uvs = faces_uvs is not None and verts_uvs is not None
+    save_texture = save_uvs and texture_map is not None
     output_path = Path(f)
 
     # Save the .obj file
@@ -778,7 +780,7 @@ def save_obj(
             faces_normals_idx=faces_normals_idx,
             verts_uvs=verts_uvs,
             faces_uvs=faces_uvs,
-            save_texture=save_texture,
+            save_uvs=save_uvs,
             save_normals=normals is not None,
         )
 
@@ -816,7 +818,7 @@ def _save(
     faces_normals_idx: Optional[torch.Tensor] = None,
     verts_uvs: Optional[torch.Tensor] = None,
     faces_uvs: Optional[torch.Tensor] = None,
-    save_texture: bool = False,
+    save_uvs: bool = False,
     save_normals: bool = False,
 ) -> None:
     if len(verts) and (verts.dim() != 2 or verts.size(1) != 3):
@@ -832,26 +834,24 @@ def _save(
         return
 
     verts, faces = verts.cpu(), faces.cpu()
+    vert_rows = verts.detach().cpu().numpy()
 
-    lines = ""
+    lines = []
 
     if decimal_places is None:
         float_str = "%f"
     else:
         float_str = "%" + ".%df" % decimal_places
 
-    if len(verts):
-        V, D = verts.shape
-        for i in range(V):
-            vert = [float_str % verts[i, j] for j in range(D)]
-            lines += "v %s\n" % " ".join(vert)
+    for vert in vert_rows:
+        lines.append("v %s\n" % " ".join(float_str % v for v in vert))
 
     if save_normals:
         assert normals is not None
         assert faces_normals_idx is not None
-        lines += _write_normals(normals, faces_normals_idx, float_str)
+        lines.append(_write_normals(normals, float_str))
 
-    if save_texture:
+    if save_uvs:
         assert faces_uvs is not None
         assert verts_uvs is not None
 
@@ -864,15 +864,13 @@ def _save(
             raise ValueError(message)
 
         verts_uvs, faces_uvs = verts_uvs.cpu(), faces_uvs.cpu()
+        uv_rows = verts_uvs.detach().cpu().numpy()
 
         # Save verts uvs after verts
-        if len(verts_uvs):
-            uV, uD = verts_uvs.shape
-            for i in range(uV):
-                uv = [float_str % verts_uvs[i, j] for j in range(uD)]
-                lines += "vt %s\n" % " ".join(uv)
+        for uv in uv_rows:
+            lines.append("vt %s\n" % " ".join(float_str % u for u in uv))
 
-    f.write(lines)
+    f.write("".join(lines))
 
     if torch.any(faces >= verts.shape[0]) or torch.any(faces < 0):
         warnings.warn("Faces have invalid indices")
@@ -881,31 +879,21 @@ def _save(
         _write_faces(
             f,
             faces,
-            faces_uvs if save_texture else None,
+            faces_uvs if save_uvs else None,
             faces_normals_idx if save_normals else None,
         )
 
 
-def _write_normals(
-    normals: torch.Tensor, faces_normals_idx: torch.Tensor, float_str: str
-) -> str:
-    if faces_normals_idx.dim() != 2 or faces_normals_idx.size(1) != 3:
-        message = (
-            "'faces_normals_idx' should either be empty or of shape (num_faces, 3)."
-        )
-        raise ValueError(message)
-
+def _write_normals(normals: torch.Tensor, float_str: str) -> str:
     if normals.dim() != 2 or normals.size(1) != 3:
         message = "'normals' should either be empty or of shape (num_verts, 3)."
         raise ValueError(message)
 
-    normals, faces_normals_idx = normals.cpu(), faces_normals_idx.cpu()
+    normal_rows = normals.detach().cpu().numpy()
 
     lines = []
-    V, D = normals.shape
-    for i in range(V):
-        normal = [float_str % normals[i, j] for j in range(D)]
-        lines.append("vn %s\n" % " ".join(normal))
+    for normal in normal_rows:
+        lines.append("vn %s\n" % " ".join(float_str % n for n in normal))
     return "".join(lines)
 
 
@@ -915,31 +903,37 @@ def _write_faces(
     faces_uvs: Optional[torch.Tensor],
     faces_normals_idx: Optional[torch.Tensor],
 ) -> None:
+    face_rows = faces.detach().cpu().numpy()
+    uv_rows = faces_uvs.detach().cpu().numpy() if faces_uvs is not None else None
+    normal_rows = (
+        faces_normals_idx.detach().cpu().numpy()
+        if faces_normals_idx is not None
+        else None
+    )
+
     F, P = faces.shape
     for i in range(F):
-        if faces_normals_idx is not None:
-            if faces_uvs is not None:
+        face_row = face_rows[i]
+        if normal_rows is not None:
+            normal_row = normal_rows[i]
+            if uv_rows is not None:
+                uv_row = uv_rows[i]
                 # Format faces as {verts_idx}/{verts_uvs_idx}/{verts_normals_idx}
                 face = [
-                    "%d/%d/%d"
-                    % (
-                        faces[i, j] + 1,
-                        faces_uvs[i, j] + 1,
-                        faces_normals_idx[i, j] + 1,
-                    )
+                    "%d/%d/%d" % (face_row[j] + 1, uv_row[j] + 1, normal_row[j] + 1)
                     for j in range(P)
                 ]
             else:
                 # Format faces as {verts_idx}//{verts_normals_idx}
                 face = [
-                    "%d//%d" % (faces[i, j] + 1, faces_normals_idx[i, j] + 1)
-                    for j in range(P)
+                    "%d//%d" % (face_row[j] + 1, normal_row[j] + 1) for j in range(P)
                 ]
-        elif faces_uvs is not None:
+        elif uv_rows is not None:
+            uv_row = uv_rows[i]
             # Format faces as {verts_idx}/{verts_uvs_idx}
-            face = ["%d/%d" % (faces[i, j] + 1, faces_uvs[i, j] + 1) for j in range(P)]
+            face = ["%d/%d" % (face_row[j] + 1, uv_row[j] + 1) for j in range(P)]
         else:
-            face = ["%d" % (faces[i, j] + 1) for j in range(P)]
+            face = ["%d" % (face_row[j] + 1) for j in range(P)]
 
         if i + 1 < F:
             f.write("f %s\n" % " ".join(face))

tests/test_io_obj.py

@@ -955,6 +955,60 @@ class TestMeshObjIO(TestCaseMixin, unittest.TestCase):
             with open(obj_file, "r") as actual_file:
                 self.assertEqual(actual_file.read(), expected_obj_file)
 
+    def test_save_obj_with_uvs_without_texture_map(self):
+        verts = torch.tensor(
+            [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0]],
+            dtype=torch.float32,
+        )
+        faces = torch.tensor([[0, 2, 1]], dtype=torch.int64)
+        verts_uvs = torch.tensor(
+            [[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]],
+            dtype=torch.float32,
+        )
+        faces_uvs = torch.tensor([[3, 1, 2]], dtype=torch.int64)
+        normals = torch.tensor([[0.0, 0.0, 1.0], [0.0, 1.0, 0.0]])
+        faces_normals_idx = torch.tensor([[1, 0, 1]], dtype=torch.int64)
+
+        for save_normals in (False, True):
+            with (
+                self.subTest(save_normals=save_normals),
+                TemporaryDirectory() as temp_dir,
+            ):
+                obj_file = Path(temp_dir) / "mesh.obj"
+                save_obj(
+                    obj_file,
+                    verts,
+                    faces,
+                    verts_uvs=verts_uvs,
+                    faces_uvs=faces_uvs,
+                    normals=normals if save_normals else None,
+                    faces_normals_idx=faces_normals_idx if save_normals else None,
+                )
+
+                self.assertEqual(os.listdir(temp_dir), ["mesh.obj"])
+                obj_text = obj_file.read_text()
+                self.assertNotIn("mtllib", obj_text)
+                self.assertNotIn("usemtl", obj_text)
+                expected_face = (
+                    "f 1/4/2 3/2/1 2/3/2" if save_normals else "f 1/4 3/2 2/3"
+                )
+                self.assertIn(expected_face, obj_text)
+
+                loaded_verts, loaded_faces, aux = load_obj(
+                    obj_file, load_textures=False
+                )
+                self.assertClose(loaded_verts, verts)
+                self.assertClose(loaded_faces.verts_idx, faces)
+                self.assertClose(aux.verts_uvs, verts_uvs)
+                self.assertClose(loaded_faces.textures_idx, faces_uvs)
+                self.assertIsNone(aux.material_colors)
+                self.assertIsNone(aux.texture_images)
+                if save_normals:
+                    self.assertClose(aux.normals, normals)
+                    self.assertClose(loaded_faces.normals_idx, faces_normals_idx)
+                else:
+                    self.assertIsNone(aux.normals)
+
     def test_save_obj_with_texture(self):
         verts = torch.tensor(
             [[0.01, 0.2, 0.301], [0.2, 0.03, 0.408], [0.3, 0.4, 0.05], [0.6, 0.7, 0.8]],

tests/test_symeig3x3.py

@@ -66,6 +66,51 @@ class TestSymEig3x3(TestCaseMixin, unittest.TestCase):
             rtol=rtol,
         )
 
+    def test_is_eigen_scaled_gpu(self):
+        self._test_is_eigen_scaled(self._gpu)
+
+    def test_is_eigen_scaled_cpu(self):
+        self._test_is_eigen_scaled(self._cpu)
+
+    def _test_is_eigen_scaled(self, device):
+        """
+        The decomposition is homogeneous, so scaling the input must not affect
+        the relative accuracy of the result.
+        """
+        test_input = self.create_random_sym3x3(device, n=self.TEST_BATCH_SIZE)
+
+        for scale in (1e-6, 1e-4, 1e-2, 1e2, 1e4, 1e6):
+            scaled_input = test_input * scale
+
+            self._test_is_not_nan_or_inf(scaled_input)
+            self._test_is_eigen(scaled_input, atol=1e-04 * scale, rtol=1e-02)
+            self._test_eigenvectors_are_orthonormal(scaled_input)
+
+    def test_mixed_magnitudes_gpu(self):
+        self._test_mixed_magnitudes(self._gpu)
+
+    def test_mixed_magnitudes_cpu(self):
+        self._test_mixed_magnitudes(self._cpu)
+
+    def _test_mixed_magnitudes(self, device):
+        """
+        Each matrix in a batch must be conditioned on its own magnitude, so
+        matrices spanning many orders of magnitude can be batched together.
+        """
+        test_input = self.create_random_sym3x3(device, n=self.TEST_BATCH_SIZE)
+        scales = torch.logspace(-6, 6, self.TEST_BATCH_SIZE, device=device)[
+            :, None, None
+        ]
+        scaled_input = test_input * scales
+
+        eigenvalues, eigenvectors = symeig3x3(scaled_input, eigenvectors=True)
+        reconstruction = (
+            eigenvectors @ eigenvalues.diag_embed() @ eigenvectors.transpose(-2, -1)
+        )
+
+        self.assertClose(test_input, reconstruction / scales, atol=1e-04, rtol=1e-02)
+        self._test_eigenvectors_are_orthonormal(scaled_input)
+
     def test_eigenvectors_are_orthonormal_gpu(self):
         test_input = self.create_random_sym3x3(self._gpu, n=self.TEST_BATCH_SIZE)
 
@@ -192,21 +237,31 @@ class TestSymEig3x3(TestCaseMixin, unittest.TestCase):
     def test_degenerate_eigenvalues_cpu(self):
         self._test_degenerate_eigenvalues(self._cpu)
 
-    def _test_degenerate_eigenvalues(self, device):
+    def test_degenerate_eigenvalues_scaled_gpu(self):
+        for scale in (1e-6, 1e6):
+            self._test_degenerate_eigenvalues(self._gpu, scale=scale)
+
+    def test_degenerate_eigenvalues_scaled_cpu(self):
+        for scale in (1e-6, 1e6):
+            self._test_degenerate_eigenvalues(self._cpu, scale=scale)
+
+    def _test_degenerate_eigenvalues(self, device, scale: float = 1.0):
         """
-        Test degenerate eigenvalues like zero-valued and with 2-/3-multiplicity
+        Test degenerate eigenvalues like zero-valued and with 2-/3-multiplicity.
+        These are the cases held together by the regularizers inside the
+        decomposition, which only bite at the right input magnitude.
         """
         # Error tolerances for degenerate values are increased as things might become
         #  numerically unstable
-        deg_atol = 1e-3
+        deg_atol = 1e-3 * scale
         deg_rtol = 1.0
 
         # Construct random orthonormal sets
         test_eigenvecs = random_rotations(n=self.TEST_BATCH_SIZE, device=device)
 
         # Construct random eigenvalues
-        test_eigenvals = torch.randn(
-            (self.TEST_BATCH_SIZE, 3), device=test_eigenvecs.device
+        test_eigenvals = (
+            torch.randn((self.TEST_BATCH_SIZE, 3), device=test_eigenvecs.device) * scale
         )
         self._test_eigenvalues_and_eigenvectors(
             test_eigenvecs, test_eigenvals, atol=deg_atol, rtol=deg_rtol