From 860b742a02e62ec85f48929268349916ca4ce8a5 Mon Sep 17 00:00:00 2001
From: Jeremy Reizenstein <reizenstein@fb.com>
Date: Thu, 23 Sep 2021 06:57:11 -0700
Subject: [PATCH] deterministic rasterization

Summary: Attempt to fix #659, an observation that the rasterizer is nondeterministic, by resolving tied faces by picking those with lower index.

Reviewed By: nikhilaravi, patricklabatut

Differential Revision: D30699039

fbshipit-source-id: 39ed797eb7e9ce7370ae71259ad6b757f9449923
---
 .../csrc/rasterize_meshes/rasterize_meshes.cu |  2 +-
 .../rasterize_meshes/rasterize_meshes_cpu.cpp |  9 +-------
 tests/test_rasterize_meshes.py                | 22 +++++++++++++++++++
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/pytorch3d/csrc/rasterize_meshes/rasterize_meshes.cu b/pytorch3d/csrc/rasterize_meshes/rasterize_meshes.cu
index 73e899741..c044e58d9 100644
--- a/pytorch3d/csrc/rasterize_meshes/rasterize_meshes.cu
+++ b/pytorch3d/csrc/rasterize_meshes/rasterize_meshes.cu
@@ -28,7 +28,7 @@ struct Pixel {
 };
 
 __device__ bool operator<(const Pixel& a, const Pixel& b) {
-  return a.z < b.z;
+  return a.z < b.z || (a.z == b.z && a.idx < b.idx);
 }
 
 // Get the xyz coordinates of the three vertices for the face given by the
diff --git a/pytorch3d/csrc/rasterize_meshes/rasterize_meshes_cpu.cpp b/pytorch3d/csrc/rasterize_meshes/rasterize_meshes_cpu.cpp
index 37f031052..8b2a85e6d 100644
--- a/pytorch3d/csrc/rasterize_meshes/rasterize_meshes_cpu.cpp
+++ b/pytorch3d/csrc/rasterize_meshes/rasterize_meshes_cpu.cpp
@@ -117,13 +117,6 @@ struct IsNeighbor {
   int neighbor_idx;
 };
 
-// Function to sort based on the z distance in the top K queue
-bool SortTopKByZdist(
-    std::tuple<float, int, float, float, float, float> a,
-    std::tuple<float, int, float, float, float, float> b) {
-  return std::get<0>(a) < std::get<0>(b);
-}
-
 std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
 RasterizeMeshesNaiveCpu(
     const torch::Tensor& face_verts,
@@ -310,7 +303,7 @@ RasterizeMeshesNaiveCpu(
 
           // Sort the deque inplace based on the z distance
           // to mimic using a priority queue.
-          std::sort(q.begin(), q.end(), SortTopKByZdist);
+          std::sort(q.begin(), q.end());
           if (static_cast<int>(q.size()) > K) {
             // remove the last value
             q.pop_back();
diff --git a/tests/test_rasterize_meshes.py b/tests/test_rasterize_meshes.py
index 306c22150..033370bb8 100644
--- a/tests/test_rasterize_meshes.py
+++ b/tests/test_rasterize_meshes.py
@@ -1151,6 +1151,28 @@ def _test_coarse_rasterize(self, device):
         bin_faces_same = (bin_faces.squeeze() == bin_faces_expected).all()
         self.assertTrue(bin_faces_same.item() == 1)
 
+    def test_order_of_ties(self):
+        # Tied faces are rasterized in index order
+        # We rasterize a mesh with many faces.
+        device = torch.device("cuda:0")
+        verts = -5 * torch.eye(3, dtype=torch.float32, device=device)[None]
+        faces = torch.arange(3, device=device, dtype=torch.int64).expand(1, 100, 3)
+        mesh = Meshes(verts=verts, faces=faces)
+
+        R, T = look_at_view_transform(2.7, 0.0, 0.0)
+        cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
+
+        raster_settings = RasterizationSettings(
+            image_size=28, faces_per_pixel=100, bin_size=0
+        )
+        rasterizer = MeshRasterizer(raster_settings=raster_settings)
+
+        out = rasterizer(mesh, cameras=cameras)
+        self.assertClose(
+            out.pix_to_face[0, 14:, :14],
+            torch.arange(100, device=device).expand(14, 14, 100),
+        )
+
     @staticmethod
     def rasterize_meshes_python_with_init(
         num_meshes: int,