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/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