CPU implementation for point_mesh functions

Summary:
point_mesh functions were missing CPU implementations.
The indices returned are not always matching, possibly due to numerical instability.

Reviewed By: gkioxari

Differential Revision: D21594264

fbshipit-source-id: 3016930e2a9a0f3cd8b3ac4c94a92c9411c0989d

pytorch3d/csrc/point_mesh/point_mesh.cpp

+// Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+
+#include <torch/extension.h>
+#include <array>
+#include <limits>
+#include "utils/geometry_utils.h"
+#include "utils/vec3.h"
+
+// - We start with implementations of simple operations on points, edges and
+// faces. The hull of H points is a point if H=1, an edge if H=2, a face if H=3.
+
+template <typename T>
+vec3<T> ExtractPoint(const at::TensorAccessor<T, 1>& t) {
+  return vec3(t[0], t[1], t[2]);
+}
+
+template <class Accessor>
+struct ExtractHullHelper {
+  template <int H>
+  static std::array<vec3<std::remove_pointer_t<typename Accessor::PtrType>>, H>
+  get(const Accessor& t);
+
+  template <>
+  static std::array<vec3<std::remove_pointer_t<typename Accessor::PtrType>>, 1>
+  get<1>(const Accessor& t) {
+    return {ExtractPoint(t)};
+  }
+
+  template <>
+  static std::array<vec3<std::remove_pointer_t<typename Accessor::PtrType>>, 2>
+  get<2>(const Accessor& t) {
+    return {ExtractPoint(t[0]), ExtractPoint(t[1])};
+  }
+
+  template <>
+  static std::array<vec3<std::remove_pointer_t<typename Accessor::PtrType>>, 3>
+  get<3>(const Accessor& t) {
+    return {ExtractPoint(t[0]), ExtractPoint(t[1]), ExtractPoint(t[2])};
+  }
+};
+
+template <int H, typename Accessor>
+std::array<vec3<std::remove_pointer_t<typename Accessor::PtrType>>, H>
+ExtractHull(const Accessor& t) {
+  return ExtractHullHelper<Accessor>::template get<H>(t);
+}
+
+template <typename T>
+void IncrementPoint(at::TensorAccessor<T, 1>&& t, const vec3<T>& point) {
+  t[0] += point.x;
+  t[1] += point.y;
+  t[2] += point.z;
+}
+
+// distance between the convex hull of A points and B points
+// this could be done in c++17 with tuple_cat and invoke
+template <typename T>
+T HullDistance(
+    const std::array<vec3<T>, 1>& a,
+    const std::array<vec3<T>, 2>& b) {
+  using std::get;
+  return PointLine3DistanceForward(get<0>(a), get<0>(b), get<1>(b));
+}
+template <typename T>
+T HullDistance(
+    const std::array<vec3<T>, 1>& a,
+    const std::array<vec3<T>, 3>& b) {
+  using std::get;
+  return PointTriangle3DistanceForward(
+      get<0>(a), get<0>(b), get<1>(b), get<2>(b));
+}
+template <typename T>
+T HullDistance(
+    const std::array<vec3<T>, 2>& a,
+    const std::array<vec3<T>, 1>& b) {
+  return HullDistance(b, a);
+}
+template <typename T>
+T HullDistance(
+    const std::array<vec3<T>, 3>& a,
+    const std::array<vec3<T>, 1>& b) {
+  return HullDistance(b, a);
+}
+
+template <typename T>
+void HullHullDistanceBackward(
+    const std::array<vec3<T>, 1>& a,
+    const std::array<vec3<T>, 2>& b,
+    T grad_dist,
+    at::TensorAccessor<T, 1>&& grad_a,
+    at::TensorAccessor<T, 2>&& grad_b) {
+  using std::get;
+  auto res =
+      PointLine3DistanceBackward(get<0>(a), get<0>(b), get<1>(b), grad_dist);
+  IncrementPoint(std::move(grad_a), get<0>(res));
+  IncrementPoint(grad_b[0], get<1>(res));
+  IncrementPoint(grad_b[1], get<2>(res));
+}
+template <typename T>
+void HullHullDistanceBackward(
+    const std::array<vec3<T>, 1>& a,
+    const std::array<vec3<T>, 3>& b,
+    T grad_dist,
+    at::TensorAccessor<T, 1>&& grad_a,
+    at::TensorAccessor<T, 2>&& grad_b) {
+  using std::get;
+  auto res = PointTriangle3DistanceBackward(
+      get<0>(a), get<0>(b), get<1>(b), get<2>(b), grad_dist);
+  IncrementPoint(std::move(grad_a), get<0>(res));
+  IncrementPoint(grad_b[0], get<1>(res));
+  IncrementPoint(grad_b[1], get<2>(res));
+  IncrementPoint(grad_b[2], get<3>(res));
+}
+template <typename T>
+void HullHullDistanceBackward(
+    const std::array<vec3<T>, 3>& a,
+    const std::array<vec3<T>, 1>& b,
+    T grad_dist,
+    at::TensorAccessor<T, 2>&& grad_a,
+    at::TensorAccessor<T, 1>&& grad_b) {
+  return HullHullDistanceBackward(
+      b, a, grad_dist, std::move(grad_b), std::move(grad_a));
+}
+template <typename T>
+void HullHullDistanceBackward(
+    const std::array<vec3<T>, 2>& a,
+    const std::array<vec3<T>, 1>& b,
+    T grad_dist,
+    at::TensorAccessor<T, 2>&& grad_a,
+    at::TensorAccessor<T, 1>&& grad_b) {
+  return HullHullDistanceBackward(
+      b, a, grad_dist, std::move(grad_b), std::move(grad_a));
+}
+
+template <int H>
+void ValidateShape(const at::Tensor& as) {
+  if (H == 1) {
+    TORCH_CHECK(as.size(1) == 3);
+  } else {
+    TORCH_CHECK(as.size(2) == 3 && as.size(1) == H);
+  }
+}
+
+// ----------- Here begins the implementation of each top-level
+//             function using non-type template parameters to
+//             implement all the cases in one go. ----------- //
+
+template <int H1, int H2>
+std::tuple<at::Tensor, at::Tensor> HullHullDistanceForwardCpu(
+    const at::Tensor& as,
+    const at::Tensor& as_first_idx,
+    const at::Tensor& bs,
+    const at::Tensor& bs_first_idx) {
+  const int64_t A_N = as.size(0);
+  const int64_t B_N = bs.size(0);
+  const int64_t BATCHES = as_first_idx.size(0);
+
+  ValidateShape<H1>(as);
+  ValidateShape<H2>(bs);
+
+  TORCH_CHECK(bs_first_idx.size(0) == BATCHES);
+
+  // clang-format off
+  at::Tensor dists = at::zeros({A_N,}, as.options());
+  at::Tensor idxs = at::zeros({A_N,}, as_first_idx.options());
+  // clang-format on
+
+  auto as_a = as.accessor < float, H1 == 1 ? 2 : 3 > ();
+  auto bs_a = bs.accessor < float, H2 == 1 ? 2 : 3 > ();
+  auto as_first_idx_a = as_first_idx.accessor<int64_t, 1>();
+  auto bs_first_idx_a = bs_first_idx.accessor<int64_t, 1>();
+  auto dists_a = dists.accessor<float, 1>();
+  auto idxs_a = idxs.accessor<int64_t, 1>();
+  int64_t a_batch_end = 0;
+  int64_t b_batch_start = 0, b_batch_end = 0;
+  int64_t batch_idx = 0;
+  for (int64_t a_n = 0; a_n < A_N; ++a_n) {
+    if (a_n == a_batch_end) {
+      ++batch_idx;
+      b_batch_start = b_batch_end;
+      if (batch_idx == BATCHES) {
+        a_batch_end = std::numeric_limits<int64_t>::max();
+        b_batch_end = B_N;
+      } else {
+        a_batch_end = as_first_idx_a[batch_idx];
+        b_batch_end = bs_first_idx_a[batch_idx];
+      }
+    }
+    float min_dist = std::numeric_limits<float>::max();
+    size_t min_idx = 0;
+    auto a = ExtractHull<H1>(as_a[a_n]);
+    for (int64_t b_n = b_batch_start; b_n < b_batch_end; ++b_n) {
+      float dist = HullDistance(a, ExtractHull<H2>(bs_a[b_n]));
+      if (dist <= min_dist) {
+        min_dist = dist;
+        min_idx = b_n;
+      }
+    }
+    dists_a[a_n] = min_dist;
+    idxs_a[a_n] = min_idx;
+  }
+
+  return std::make_tuple(dists, idxs);
+}
+
+template <int H1, int H2>
+std::tuple<at::Tensor, at::Tensor> HullHullDistanceBackwardCpu(
+    const at::Tensor& as,
+    const at::Tensor& bs,
+    const at::Tensor& idx_bs,
+    const at::Tensor& grad_dists) {
+  const int64_t A_N = as.size(0);
+
+  TORCH_CHECK(idx_bs.size(0) == A_N);
+  TORCH_CHECK(grad_dists.size(0) == A_N);
+  ValidateShape<H1>(as);
+  ValidateShape<H2>(bs);
+
+  at::Tensor grad_as = at::zeros_like(as);
+  at::Tensor grad_bs = at::zeros_like(bs);
+
+  auto as_a = as.accessor < float, H1 == 1 ? 2 : 3 > ();
+  auto bs_a = bs.accessor < float, H2 == 1 ? 2 : 3 > ();
+  auto grad_as_a = grad_as.accessor < float, H1 == 1 ? 2 : 3 > ();
+  auto grad_bs_a = grad_bs.accessor < float, H2 == 1 ? 2 : 3 > ();
+  auto idx_bs_a = idx_bs.accessor<int64_t, 1>();
+  auto grad_dists_a = grad_dists.accessor<float, 1>();
+
+  for (int64_t a_n = 0; a_n < A_N; ++a_n) {
+    auto a = ExtractHull<H1>(as_a[a_n]);
+    auto b = ExtractHull<H2>(bs_a[idx_bs_a[a_n]]);
+    HullHullDistanceBackward(
+        a, b, grad_dists_a[a_n], grad_as_a[a_n], grad_bs_a[idx_bs_a[a_n]]);
+  }
+  return std::make_tuple(grad_as, grad_bs);
+}
+
+template <int H>
+torch::Tensor PointHullArrayDistanceForwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& bs) {
+  const int64_t P = points.size(0);
+  const int64_t B_N = bs.size(0);
+
+  TORCH_CHECK(points.size(1) == 3, "points must be of shape Px3");
+  ValidateShape<H>(bs);
+
+  at::Tensor dists = at::zeros({P, B_N}, points.options());
+  auto points_a = points.accessor<float, 2>();
+  auto bs_a = bs.accessor<float, 3>();
+  auto dists_a = dists.accessor<float, 2>();
+  for (int64_t p = 0; p < P; ++p) {
+    auto point = ExtractHull<1>(points_a[p]);
+    auto dest = dists_a[p];
+    for (int64_t b_n = 0; b_n < B_N; ++b_n) {
+      auto b = ExtractHull<H>(bs_a[b_n]);
+      dest[b_n] = HullDistance(point, b);
+    }
+  }
+  return dists;
+}
+
+template <int H>
+std::tuple<at::Tensor, at::Tensor> PointHullArrayDistanceBackwardCpu(
+    const at::Tensor& points,
+    const at::Tensor& bs,
+    const at::Tensor& grad_dists) {
+  const int64_t P = points.size(0);
+  const int64_t B_N = bs.size(0);
+
+  TORCH_CHECK(points.size(1) == 3, "points must be of shape Px3");
+  ValidateShape<H>(bs);
+  TORCH_CHECK((grad_dists.size(0) == P) && (grad_dists.size(1) == B_N));
+
+  at::Tensor grad_points = at::zeros({P, 3}, points.options());
+  at::Tensor grad_bs = at::zeros({B_N, H, 3}, bs.options());
+
+  auto points_a = points.accessor<float, 2>();
+  auto bs_a = bs.accessor<float, 3>();
+  auto grad_dists_a = grad_dists.accessor<float, 2>();
+  auto grad_points_a = grad_points.accessor<float, 2>();
+  auto grad_bs_a = grad_bs.accessor<float, 3>();
+  for (int64_t p = 0; p < P; ++p) {
+    auto point = ExtractHull<1>(points_a[p]);
+    auto grad_point = grad_points_a[p];
+    auto grad_dist = grad_dists_a[p];
+    for (int64_t b_n = 0; b_n < B_N; ++b_n) {
+      auto b = ExtractHull<H>(bs_a[b_n]);
+      HullHullDistanceBackward(
+          point, b, grad_dist[b_n], std::move(grad_point), grad_bs_a[b_n]);
+    }
+  }
+  return std::make_tuple(grad_points, grad_bs);
+}
+
+// ---------- Here begin the exported functions ------------ //
+
+std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceForwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& points_first_idx,
+    const torch::Tensor& tris,
+    const torch::Tensor& tris_first_idx) {
+  return HullHullDistanceForwardCpu<1, 3>(
+      points, points_first_idx, tris, tris_first_idx);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceBackwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& tris,
+    const torch::Tensor& idx_points,
+    const torch::Tensor& grad_dists) {
+  return HullHullDistanceBackwardCpu<1, 3>(
+      points, tris, idx_points, grad_dists);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceForwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& points_first_idx,
+    const torch::Tensor& tris,
+    const torch::Tensor& tris_first_idx) {
+  return HullHullDistanceForwardCpu<3, 1>(
+      tris, tris_first_idx, points, points_first_idx);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceBackwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& tris,
+    const torch::Tensor& idx_tris,
+    const torch::Tensor& grad_dists) {
+  auto res =
+      HullHullDistanceBackwardCpu<3, 1>(tris, points, idx_tris, grad_dists);
+  return std::make_tuple(std::get<1>(res), std::get<0>(res));
+}
+
+torch::Tensor PointEdgeArrayDistanceForwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& segms) {
+  return PointHullArrayDistanceForwardCpu<2>(points, segms);
+}
+
+std::tuple<at::Tensor, at::Tensor> PointFaceArrayDistanceBackwardCpu(
+    const at::Tensor& points,
+    const at::Tensor& tris,
+    const at::Tensor& grad_dists) {
+  return PointHullArrayDistanceBackwardCpu<3>(points, tris, grad_dists);
+}
+
+torch::Tensor PointFaceArrayDistanceForwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& tris) {
+  return PointHullArrayDistanceForwardCpu<3>(points, tris);
+}
+
+std::tuple<at::Tensor, at::Tensor> PointEdgeArrayDistanceBackwardCpu(
+    const at::Tensor& points,
+    const at::Tensor& segms,
+    const at::Tensor& grad_dists) {
+  return PointHullArrayDistanceBackwardCpu<2>(points, segms, grad_dists);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> PointEdgeDistanceForwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& points_first_idx,
+    const torch::Tensor& segms,
+    const torch::Tensor& segms_first_idx,
+    const int64_t /*max_points*/) {
+  return HullHullDistanceForwardCpu<1, 2>(
+      points, points_first_idx, segms, segms_first_idx);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> PointEdgeDistanceBackwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& segms,
+    const torch::Tensor& idx_points,
+    const torch::Tensor& grad_dists) {
+  return HullHullDistanceBackwardCpu<1, 2>(
+      points, segms, idx_points, grad_dists);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> EdgePointDistanceForwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& points_first_idx,
+    const torch::Tensor& segms,
+    const torch::Tensor& segms_first_idx,
+    const int64_t /*max_segms*/) {
+  return HullHullDistanceForwardCpu<2, 1>(
+      segms, segms_first_idx, points, points_first_idx);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> EdgePointDistanceBackwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& segms,
+    const torch::Tensor& idx_segms,
+    const torch::Tensor& grad_dists) {
+  auto res =
+      HullHullDistanceBackwardCpu<2, 1>(segms, points, idx_segms, grad_dists);
+  return std::make_tuple(std::get<1>(res), std::get<0>(res));
+}

pytorch3d/csrc/point_mesh/point_mesh_edge.h

@@ -46,6 +46,13 @@ std::tuple<torch::Tensor, torch::Tensor> PointEdgeDistanceForwardCuda(
     const int64_t max_points);
 #endif
 
+std::tuple<torch::Tensor, torch::Tensor> PointEdgeDistanceForwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& points_first_idx,
+    const torch::Tensor& segms,
+    const torch::Tensor& segms_first_idx,
+    const int64_t max_points);
+
 std::tuple<torch::Tensor, torch::Tensor> PointEdgeDistanceForward(
     const torch::Tensor& points,
     const torch::Tensor& points_first_idx,
@@ -64,7 +71,8 @@ std::tuple<torch::Tensor, torch::Tensor> PointEdgeDistanceForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
-  AT_ERROR("No CPU implementation.");
+  return PointEdgeDistanceForwardCpu(
+      points, points_first_idx, segms, segms_first_idx, max_points);
 }
 
 // Backward pass for PointEdgeDistance.
@@ -91,6 +99,12 @@ std::tuple<torch::Tensor, torch::Tensor> PointEdgeDistanceBackwardCuda(
     const torch::Tensor& grad_dists);
 #endif
 
+std::tuple<torch::Tensor, torch::Tensor> PointEdgeDistanceBackwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& segms,
+    const torch::Tensor& idx_points,
+    const torch::Tensor& grad_dists);
+
 std::tuple<torch::Tensor, torch::Tensor> PointEdgeDistanceBackward(
     const torch::Tensor& points,
     const torch::Tensor& segms,
@@ -107,7 +121,7 @@ std::tuple<torch::Tensor, torch::Tensor> PointEdgeDistanceBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
-  AT_ERROR("No CPU implementation.");
+  return PointEdgeDistanceBackwardCpu(points, segms, idx_points, grad_dists);
 }
 
 // ****************************************************************************
@@ -150,6 +164,13 @@ std::tuple<torch::Tensor, torch::Tensor> EdgePointDistanceForwardCuda(
     const int64_t max_segms);
 #endif
 
+std::tuple<torch::Tensor, torch::Tensor> EdgePointDistanceForwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& points_first_idx,
+    const torch::Tensor& segms,
+    const torch::Tensor& segms_first_idx,
+    const int64_t max_segms);
+
 std::tuple<torch::Tensor, torch::Tensor> EdgePointDistanceForward(
     const torch::Tensor& points,
     const torch::Tensor& points_first_idx,
@@ -168,7 +189,8 @@ std::tuple<torch::Tensor, torch::Tensor> EdgePointDistanceForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
-  AT_ERROR("No CPU implementation.");
+  return EdgePointDistanceForwardCpu(
+      points, points_first_idx, segms, segms_first_idx, max_segms);
 }
 
 // Backward pass for EdgePointDistance.
@@ -195,6 +217,12 @@ std::tuple<torch::Tensor, torch::Tensor> EdgePointDistanceBackwardCuda(
     const torch::Tensor& grad_dists);
 #endif
 
+std::tuple<torch::Tensor, torch::Tensor> EdgePointDistanceBackwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& segms,
+    const torch::Tensor& idx_segms,
+    const torch::Tensor& grad_dists);
+
 std::tuple<torch::Tensor, torch::Tensor> EdgePointDistanceBackward(
     const torch::Tensor& points,
     const torch::Tensor& segms,
@@ -211,7 +239,7 @@ std::tuple<torch::Tensor, torch::Tensor> EdgePointDistanceBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
-  AT_ERROR("No CPU implementation.");
+  return EdgePointDistanceBackwardCpu(points, segms, idx_segms, grad_dists);
 }
 
 // ****************************************************************************
@@ -242,6 +270,10 @@ torch::Tensor PointEdgeArrayDistanceForwardCuda(
     const torch::Tensor& segms);
 #endif
 
+torch::Tensor PointEdgeArrayDistanceForwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& segms);
+
 torch::Tensor PointEdgeArrayDistanceForward(
     const torch::Tensor& points,
     const torch::Tensor& segms) {
@@ -254,7 +286,7 @@ torch::Tensor PointEdgeArrayDistanceForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
-  AT_ERROR("No CPU implementation.");
+  return PointEdgeArrayDistanceForwardCpu(points, segms);
 }
 
 // Backward pass for PointEdgeArrayDistance.
@@ -277,6 +309,11 @@ std::tuple<torch::Tensor, torch::Tensor> PointEdgeArrayDistanceBackwardCuda(
     const torch::Tensor& grad_dists);
 #endif
 
+std::tuple<torch::Tensor, torch::Tensor> PointEdgeArrayDistanceBackwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& segms,
+    const torch::Tensor& grad_dists);
+
 std::tuple<torch::Tensor, torch::Tensor> PointEdgeArrayDistanceBackward(
     const torch::Tensor& points,
     const torch::Tensor& segms,
@@ -291,5 +328,5 @@ std::tuple<torch::Tensor, torch::Tensor> PointEdgeArrayDistanceBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
-  AT_ERROR("No CPU implementation.");
+  return PointEdgeArrayDistanceBackwardCpu(points, segms, grad_dists);
 }

pytorch3d/csrc/point_mesh/point_mesh_face.h

@@ -19,7 +19,7 @@
 //    points_first_idx: LongTensor of shape (N,) indicating the first point
 //        index for each example in the batch
 //    tris: FloatTensor of shape (T, 3, 3) of the triangular faces. The t-th
-//        triangulare face is spanned by (tris[t, 0], tris[t, 1], tris[t, 2])
+//        triangular face is spanned by (tris[t, 0], tris[t, 1], tris[t, 2])
 //    tris_first_idx: LongTensor of shape (N,) indicating the first face
 //        index for each example in the batch
 //    max_points: Scalar equal to max(P_i) for i in [0, N - 1] containing
@@ -48,6 +48,12 @@ std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceForwardCuda(
     const int64_t max_points);
 #endif
 
+std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceForwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& points_first_idx,
+    const torch::Tensor& tris,
+    const torch::Tensor& tris_first_idx);
+
 std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceForward(
     const torch::Tensor& points,
     const torch::Tensor& points_first_idx,
@@ -66,7 +72,8 @@ std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
-  AT_ERROR("No CPU implementation.");
+  return PointFaceDistanceForwardCpu(
+      points, points_first_idx, tris, tris_first_idx);
 }
 
 // Backward pass for PointFaceDistance.
@@ -92,6 +99,11 @@ std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceBackwardCuda(
     const torch::Tensor& idx_points,
     const torch::Tensor& grad_dists);
 #endif
+std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceBackwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& tris,
+    const torch::Tensor& idx_points,
+    const torch::Tensor& grad_dists);
 
 std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceBackward(
     const torch::Tensor& points,
@@ -109,7 +121,7 @@ std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
-  AT_ERROR("No CPU implementation.");
+  return PointFaceDistanceBackwardCpu(points, tris, idx_points, grad_dists);
 }
 
 // ****************************************************************************
@@ -124,7 +136,7 @@ std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceBackward(
 //    points_first_idx: LongTensor of shape (N,) indicating the first point
 //        index for each example in the batch
 //    tris: FloatTensor of shape (T, 3, 3) of the triangular faces. The t-th
-//        triangulare face is spanned by (tris[t, 0], tris[t, 1], tris[t, 2])
+//        triangular face is spanned by (tris[t, 0], tris[t, 1], tris[t, 2])
 //    tris_first_idx: LongTensor of shape (N,) indicating the first face
 //        index for each example in the batch
 //    max_tris: Scalar equal to max(T_i) for i in [0, N - 1] containing
@@ -149,9 +161,15 @@ std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceForwardCuda(
     const torch::Tensor& points_first_idx,
     const torch::Tensor& tris,
     const torch::Tensor& tris_first_idx,
-    const int64_t max_tros);
+    const int64_t max_tris);
 #endif
 
+std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceForwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& points_first_idx,
+    const torch::Tensor& tris,
+    const torch::Tensor& tris_first_idx);
+
 std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceForward(
     const torch::Tensor& points,
     const torch::Tensor& points_first_idx,
@@ -170,7 +188,8 @@ std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
-  AT_ERROR("No CPU implementation.");
+  return FacePointDistanceForwardCpu(
+      points, points_first_idx, tris, tris_first_idx);
 }
 
 // Backward pass for FacePointDistance.
@@ -197,6 +216,12 @@ std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceBackwardCuda(
     const torch::Tensor& grad_dists);
 #endif
 
+std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceBackwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& tris,
+    const torch::Tensor& idx_tris,
+    const torch::Tensor& grad_dists);
+
 std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceBackward(
     const torch::Tensor& points,
     const torch::Tensor& tris,
@@ -213,7 +238,7 @@ std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
-  AT_ERROR("No CPU implementation.");
+  return FacePointDistanceBackwardCpu(points, tris, idx_tris, grad_dists);
 }
 
 // ****************************************************************************
@@ -226,7 +251,7 @@ std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceBackward(
 // Args:
 //    points: FloatTensor of shape (P, 3)
 //    tris: FloatTensor of shape (T, 3, 3) of the triangular faces. The t-th
-//        triangulare face is spanned by (tris[t, 0], tris[t, 1], tris[t, 2])
+//        triangular face is spanned by (tris[t, 0], tris[t, 1], tris[t, 2])
 //
 // Returns:
 //    dists: FloatTensor of shape (P, T), where dists[p, t] is the squared
@@ -245,6 +270,10 @@ torch::Tensor PointFaceArrayDistanceForwardCuda(
     const torch::Tensor& tris);
 #endif
 
+torch::Tensor PointFaceArrayDistanceForwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& tris);
+
 torch::Tensor PointFaceArrayDistanceForward(
     const torch::Tensor& points,
     const torch::Tensor& tris) {
@@ -257,7 +286,7 @@ torch::Tensor PointFaceArrayDistanceForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
-  AT_ERROR("No CPU implementation.");
+  return PointFaceArrayDistanceForwardCpu(points, tris);
 }
 
 // Backward pass for PointFaceArrayDistance.
@@ -278,6 +307,10 @@ std::tuple<torch::Tensor, torch::Tensor> PointFaceArrayDistanceBackwardCuda(
     const torch::Tensor& tris,
     const torch::Tensor& grad_dists);
 #endif
+std::tuple<torch::Tensor, torch::Tensor> PointFaceArrayDistanceBackwardCpu(
+    const torch::Tensor& points,
+    const torch::Tensor& tris,
+    const torch::Tensor& grad_dists);
 
 std::tuple<torch::Tensor, torch::Tensor> PointFaceArrayDistanceBackward(
     const torch::Tensor& points,
@@ -293,5 +326,5 @@ std::tuple<torch::Tensor, torch::Tensor> PointFaceArrayDistanceBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
-  AT_ERROR("No CPU implementation.");
+  return PointFaceArrayDistanceBackwardCpu(points, tris, grad_dists);
 }

pytorch3d/csrc/utils/geometry_utils.h

@@ -2,6 +2,7 @@
 
 #include <ATen/ATen.h>
 #include <algorithm>
+#include <tuple>
 #include <type_traits>
 #include "vec2.h"
 #include "vec3.h"
@@ -281,6 +282,23 @@ T PointLineDistanceForward(
   return dot(p - p_proj, p - p_proj);
 }
 
+template <typename T>
+T PointLine3DistanceForward(
+    const vec3<T>& p,
+    const vec3<T>& v0,
+    const vec3<T>& v1) {
+  const vec3<T> v1v0 = v1 - v0;
+  const T l2 = dot(v1v0, v1v0);
+  if (l2 <= kEpsilon) {
+    return dot(p - v1, p - v1);
+  }
+
+  const T t = dot(v1v0, p - v0) / l2;
+  const T tt = std::min(std::max(t, 0.00f), 1.00f);
+  const vec3<T> p_proj = v0 + tt * v1v0;
+  return dot(p - p_proj, p - p_proj);
+}
+
 // Backward pass for point to line distance in 2D.
 //
 // Args:
@@ -314,6 +332,51 @@ inline std::tuple<vec2<T>, vec2<T>, vec2<T>> PointLineDistanceBackward(
   return std::make_tuple(grad_p, grad_v0, grad_v1);
 }
 
+template <typename T>
+std::tuple<vec3<T>, vec3<T>, vec3<T>> PointLine3DistanceBackward(
+    const vec3<T>& p,
+    const vec3<T>& v0,
+    const vec3<T>& v1,
+    const T& grad_dist) {
+  const vec3<T> v1v0 = v1 - v0;
+  const vec3<T> pv0 = p - v0;
+  const T t_bot = dot(v1v0, v1v0);
+  const T t_top = dot(v1v0, pv0);
+
+  vec3<T> grad_p{0.0f, 0.0f, 0.0f};
+  vec3<T> grad_v0{0.0f, 0.0f, 0.0f};
+  vec3<T> grad_v1{0.0f, 0.0f, 0.0f};
+
+  const T tt = t_top / t_bot;
+
+  if (t_bot < kEpsilon) {
+    // if t_bot small, then v0 == v1,
+    // and dist = 0.5 * dot(pv0, pv0) + 0.5 * dot(pv1, pv1)
+    grad_p = grad_dist * 2.0f * pv0;
+    grad_v0 = -0.5f * grad_p;
+    grad_v1 = grad_v0;
+  } else if (tt < 0.0f) {
+    grad_p = grad_dist * 2.0f * pv0;
+    grad_v0 = -1.0f * grad_p;
+    // no gradients wrt v1
+  } else if (tt > 1.0f) {
+    grad_p = grad_dist * 2.0f * (p - v1);
+    grad_v1 = -1.0f * grad_p;
+    // no gradients wrt v0
+  } else {
+    const vec3<T> p_proj = v0 + tt * v1v0;
+    const vec3<T> diff = p - p_proj;
+    const vec3<T> grad_base = grad_dist * 2.0f * diff;
+    grad_p = grad_base - dot(grad_base, v1v0) * v1v0 / t_bot;
+    const vec3<T> dtt_v0 = (-1.0f * v1v0 - pv0 + 2.0f * tt * v1v0) / t_bot;
+    grad_v0 = (-1.0f + tt) * grad_base - dot(grad_base, v1v0) * dtt_v0;
+    const vec3<T> dtt_v1 = (pv0 - 2.0f * tt * v1v0) / t_bot;
+    grad_v1 = -dot(grad_base, v1v0) * dtt_v1 - tt * grad_base;
+  }
+
+  return std::make_tuple(grad_p, grad_v0, grad_v1);
+}
+
 // The forward pass for calculating the shortest distance between a point
 // and a triangle.
 // Ref: https://www.randygaul.net/2014/07/23/distance-point-to-line-segment/
@@ -396,3 +459,226 @@ PointTriangleDistanceBackward(
 
   return std::make_tuple(grad_p, grad_v0, grad_v1, grad_v2);
 }
+
+// Computes the squared distance of a point p relative to a triangle (v0, v1,
+// v2). If the point's projection p0 on the plane spanned by (v0, v1, v2) is
+// inside the triangle with vertices (v0, v1, v2), then the returned value is
+// the squared distance of p to its projection p0. Otherwise, the returned value
+// is the smallest squared distance of p from the line segments (v0, v1), (v0,
+// v2) and (v1, v2).
+//
+// Args:
+//     p: vec3 coordinates of a point
+//     v0, v1, v2: vec3 coordinates of the triangle vertices
+//
+// Returns:
+//     dist: Float of the squared distance
+//
+
+const float vEpsilon = 1e-8;
+
+template <typename T>
+vec3<T> BarycentricCoords3Forward(
+    const vec3<T>& p,
+    const vec3<T>& v0,
+    const vec3<T>& v1,
+    const vec3<T>& v2) {
+  vec3<T> p0 = v1 - v0;
+  vec3<T> p1 = v2 - v0;
+  vec3<T> p2 = p - v0;
+
+  const T d00 = dot(p0, p0);
+  const T d01 = dot(p0, p1);
+  const T d11 = dot(p1, p1);
+  const T d20 = dot(p2, p0);
+  const T d21 = dot(p2, p1);
+
+  const T denom = d00 * d11 - d01 * d01 + kEpsilon;
+  const T w1 = (d11 * d20 - d01 * d21) / denom;
+  const T w2 = (d00 * d21 - d01 * d20) / denom;
+  const T w0 = 1.0f - w1 - w2;
+
+  return vec3<T>(w0, w1, w2);
+}
+
+// Checks whether the point p is inside the triangle (v0, v1, v2).
+// A point is inside the triangle, if all barycentric coordinates
+// wrt the triangle are >= 0 & <= 1.
+//
+// NOTE that this function assumes that p lives on the space spanned
+// by (v0, v1, v2).
+// TODO(gkioxari) explicitly check whether p is coplanar with (v0, v1, v2)
+// and throw an error if check fails
+//
+// Args:
+//     p: vec3 coordinates of a point
+//     v0, v1, v2: vec3 coordinates of the triangle vertices
+//
+// Returns:
+//     inside: bool indicating wether p is inside triangle
+//
+template <typename T>
+static bool IsInsideTriangle(
+    const vec3<T>& p,
+    const vec3<T>& v0,
+    const vec3<T>& v1,
+    const vec3<T>& v2) {
+  vec3<T> bary = BarycentricCoords3Forward(p, v0, v1, v2);
+  bool x_in = 0.0f <= bary.x && bary.x <= 1.0f;
+  bool y_in = 0.0f <= bary.y && bary.y <= 1.0f;
+  bool z_in = 0.0f <= bary.z && bary.z <= 1.0f;
+  bool inside = x_in && y_in && z_in;
+  return inside;
+}
+
+template <typename T>
+T PointTriangle3DistanceForward(
+    const vec3<T>& p,
+    const vec3<T>& v0,
+    const vec3<T>& v1,
+    const vec3<T>& v2) {
+  vec3<T> normal = cross(v2 - v0, v1 - v0);
+  const T norm_normal = norm(normal);
+  normal = normal / (norm_normal + vEpsilon);
+
+  // p0 is the projection of p on the plane spanned by (v0, v1, v2)
+  // i.e. p0 = p + t * normal, s.t. (p0 - v0) is orthogonal to normal
+  const T t = dot(v0 - p, normal);
+  const vec3<T> p0 = p + t * normal;
+
+  bool is_inside = IsInsideTriangle(p0, v0, v1, v2);
+  T dist = 0.0f;
+
+  if ((is_inside) && (norm_normal > kEpsilon)) {
+    // if projection p0 is inside triangle spanned by (v0, v1, v2)
+    // then distance is equal to norm(p0 - p)^2
+    dist = t * t;
+  } else {
+    const float e01 = PointLine3DistanceForward(p, v0, v1);
+    const float e02 = PointLine3DistanceForward(p, v0, v2);
+    const float e12 = PointLine3DistanceForward(p, v1, v2);
+
+    dist = (e01 > e02) ? e02 : e01;
+    dist = (dist > e12) ? e12 : dist;
+  }
+
+  return dist;
+}
+
+template <typename T>
+std::tuple<vec3<T>, vec3<T>>
+cross_backward(const vec3<T>& a, const vec3<T>& b, const vec3<T>& grad_cross) {
+  const float grad_ax = -grad_cross.y * b.z + grad_cross.z * b.y;
+  const float grad_ay = grad_cross.x * b.z - grad_cross.z * b.x;
+  const float grad_az = -grad_cross.x * b.y + grad_cross.y * b.x;
+  const vec3<T> grad_a = vec3<T>(grad_ax, grad_ay, grad_az);
+
+  const float grad_bx = grad_cross.y * a.z - grad_cross.z * a.y;
+  const float grad_by = -grad_cross.x * a.z + grad_cross.z * a.x;
+  const float grad_bz = grad_cross.x * a.y - grad_cross.y * a.x;
+  const vec3<T> grad_b = vec3<T>(grad_bx, grad_by, grad_bz);
+
+  return std::make_tuple(grad_a, grad_b);
+}
+
+template <typename T>
+vec3<T> normalize_backward(const vec3<T>& a, const vec3<T>& grad_normz) {
+  const float a_norm = norm(a) + vEpsilon;
+  const vec3<T> out = a / a_norm;
+
+  const float grad_ax = grad_normz.x * (1.0f - out.x * out.x) / a_norm +
+      grad_normz.y * (-out.x * out.y) / a_norm +
+      grad_normz.z * (-out.x * out.z) / a_norm;
+  const float grad_ay = grad_normz.x * (-out.x * out.y) / a_norm +
+      grad_normz.y * (1.0f - out.y * out.y) / a_norm +
+      grad_normz.z * (-out.y * out.z) / a_norm;
+  const float grad_az = grad_normz.x * (-out.x * out.z) / a_norm +
+      grad_normz.y * (-out.y * out.z) / a_norm +
+      grad_normz.z * (1.0f - out.z * out.z) / a_norm;
+  return vec3<T>(grad_ax, grad_ay, grad_az);
+}
+
+// The backward pass for computing the squared distance of a point
+// to the triangle (v0, v1, v2).
+//
+// Args:
+//     p: xyz coordinates of a point
+//     v0, v1, v2: xyz coordinates of the triangle vertices
+//     grad_dist: Float of the gradient wrt dist
+//
+// Returns:
+//     tuple of gradients for the point and triangle:
+//        (float3 grad_p, float3 grad_v0, float3 grad_v1, float3 grad_v2)
+//
+
+template <typename T>
+static std::tuple<vec3<T>, vec3<T>, vec3<T>, vec3<T>>
+PointTriangle3DistanceBackward(
+    const vec3<T>& p,
+    const vec3<T>& v0,
+    const vec3<T>& v1,
+    const vec3<T>& v2,
+    const T& grad_dist) {
+  const vec3<T> v2v0 = v2 - v0;
+  const vec3<T> v1v0 = v1 - v0;
+  const vec3<T> v0p = v0 - p;
+  vec3<T> raw_normal = cross(v2v0, v1v0);
+  const T norm_normal = norm(raw_normal);
+  vec3<T> normal = raw_normal / (norm_normal + vEpsilon);
+
+  // p0 is the projection of p on the plane spanned by (v0, v1, v2)
+  // i.e. p0 = p + t * normal, s.t. (p0 - v0) is orthogonal to normal
+  const T t = dot(v0 - p, normal);
+  const vec3<T> p0 = p + t * normal;
+  const vec3<T> diff = t * normal;
+
+  bool is_inside = IsInsideTriangle(p0, v0, v1, v2);
+
+  vec3<T> grad_p(0.0f, 0.0f, 0.0f);
+  vec3<T> grad_v0(0.0f, 0.0f, 0.0f);
+  vec3<T> grad_v1(0.0f, 0.0f, 0.0f);
+  vec3<T> grad_v2(0.0f, 0.0f, 0.0f);
+
+  if ((is_inside) && (norm_normal > kEpsilon)) {
+    // derivative of dist wrt p
+    grad_p = -2.0f * grad_dist * t * normal;
+    // derivative of dist wrt normal
+    const vec3<T> grad_normal = 2.0f * grad_dist * t * (v0p + diff);
+    // derivative of dist wrt raw_normal
+    const vec3<T> grad_raw_normal = normalize_backward(raw_normal, grad_normal);
+    // derivative of dist wrt v2v0 and v1v0
+    const auto grad_cross = cross_backward(v2v0, v1v0, grad_raw_normal);
+    const vec3<T> grad_cross_v2v0 = std::get<0>(grad_cross);
+    const vec3<T> grad_cross_v1v0 = std::get<1>(grad_cross);
+    grad_v0 =
+        grad_dist * 2.0f * t * normal - (grad_cross_v2v0 + grad_cross_v1v0);
+    grad_v1 = grad_cross_v1v0;
+    grad_v2 = grad_cross_v2v0;
+  } else {
+    const T e01 = PointLine3DistanceForward(p, v0, v1);
+    const T e02 = PointLine3DistanceForward(p, v0, v2);
+    const T e12 = PointLine3DistanceForward(p, v1, v2);
+
+    if ((e01 <= e02) && (e01 <= e12)) {
+      // e01 is smallest
+      const auto grads = PointLine3DistanceBackward(p, v0, v1, grad_dist);
+      grad_p = std::get<0>(grads);
+      grad_v0 = std::get<1>(grads);
+      grad_v1 = std::get<2>(grads);
+    } else if ((e02 <= e01) && (e02 <= e12)) {
+      // e02 is smallest
+      const auto grads = PointLine3DistanceBackward(p, v0, v2, grad_dist);
+      grad_p = std::get<0>(grads);
+      grad_v0 = std::get<1>(grads);
+      grad_v2 = std::get<2>(grads);
+    } else if ((e12 <= e01) && (e12 <= e02)) {
+      // e12 is smallest
+      const auto grads = PointLine3DistanceBackward(p, v1, v2, grad_dist);
+      grad_p = std::get<0>(grads);
+      grad_v1 = std::get<1>(grads);
+      grad_v2 = std::get<2>(grads);
+    }
+  }
+
+  return std::make_tuple(grad_p, grad_v0, grad_v1, grad_v2);
+}

pytorch3d/csrc/utils/vec3.h

@@ -56,6 +56,11 @@ inline vec3<T> cross(const vec3<T>& a, const vec3<T>& b) {
       a.y * b.z - a.z * b.y, a.z * b.x - a.x * b.z, a.x * b.y - a.y * b.x);
 }
 
+template <typename T>
+inline T norm(const vec3<T>& a) {
+  return sqrt(dot(a, a));
+}
+
 template <typename T>
 std::ostream& operator<<(std::ostream& os, const vec3<T>& v) {
   os << "vec3(" << v.x << ", " << v.y << ", " << v.z << ")";

tests/test_point_mesh_distance.py

@@ -211,6 +211,9 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
         same = torch.rand((E,), dtype=torch.float32, device=device) > 0.5
         edges[same, 1] = edges[same, 0].clone().detach()
 
+        points_cpu = points.clone().cpu()
+        edges_cpu = edges.clone().cpu()
+
         points.requires_grad = True
         edges.requires_grad = True
         grad_dists = torch.rand((P, E), dtype=torch.float32, device=device)
@@ -224,22 +227,29 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
 
         # Cuda Forward Implementation
         dists_cuda = _C.point_edge_array_dist_forward(points, edges)
+        dists_cpu = _C.point_edge_array_dist_forward(points_cpu, edges_cpu)
 
         # Compare
         self.assertClose(dists_naive.cpu(), dists_cuda.cpu())
+        self.assertClose(dists_naive.cpu(), dists_cpu)
 
         # CUDA Bacwkard Implementation
         grad_points_cuda, grad_edges_cuda = _C.point_edge_array_dist_backward(
             points, edges, grad_dists
         )
+        grad_points_cpu, grad_edges_cpu = _C.point_edge_array_dist_backward(
+            points_cpu, edges_cpu, grad_dists.cpu()
+        )
 
         dists_naive.backward(grad_dists)
-        grad_points_naive = points.grad
-        grad_edges_naive = edges.grad
+        grad_points_naive = points.grad.cpu()
+        grad_edges_naive = edges.grad.cpu()
 
         # Compare
-        self.assertClose(grad_points_naive.cpu(), grad_points_cuda.cpu())
-        self.assertClose(grad_edges_naive.cpu(), grad_edges_cuda.cpu())
+        self.assertClose(grad_points_naive, grad_points_cuda.cpu())
+        self.assertClose(grad_edges_naive, grad_edges_cuda.cpu())
+        self.assertClose(grad_points_naive, grad_points_cpu)
+        self.assertClose(grad_edges_naive, grad_edges_cpu)
 
     def test_point_edge_distance(self):
         """
@@ -270,7 +280,7 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
             (points_packed.shape[0],), dtype=torch.float32, device=device
         )
 
-        # Cuda Implementation: forrward
+        # Cuda Implementation: forward
         dists_cuda, idx_cuda = _C.point_edge_dist_forward(
             points_packed, points_first_idx, edges_packed, edges_first_idx, max_p
         )
@@ -278,6 +288,20 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
         grad_points_cuda, grad_edges_cuda = _C.point_edge_dist_backward(
             points_packed, edges_packed, idx_cuda, grad_dists
         )
+        # Cpu Implementation: forward
+        dists_cpu, idx_cpu = _C.point_edge_dist_forward(
+            points_packed.cpu(),
+            points_first_idx.cpu(),
+            edges_packed.cpu(),
+            edges_first_idx.cpu(),
+            max_p,
+        )
+
+        # Cpu Implementation: backward
+        # Note that using idx_cpu doesn't pass - there seems to be a problem with tied results.
+        grad_points_cpu, grad_edges_cpu = _C.point_edge_dist_backward(
+            points_packed.cpu(), edges_packed.cpu(), idx_cuda.cpu(), grad_dists.cpu()
+        )
 
         # Naive Implementation: forward
         edges_list = packed_to_list(edges_packed, meshes.num_edges_per_mesh().tolist())
@@ -312,15 +336,18 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
 
         # Compare
         self.assertClose(dists_naive.cpu(), dists_cuda.cpu())
+        self.assertClose(dists_naive.cpu(), dists_cpu)
 
         # Naive Implementation: backward
         dists_naive.backward(grad_dists)
         grad_points_naive = torch.cat([cloud.grad for cloud in pcls.points_list()])
-        grad_edges_naive = edges_packed.grad
+        grad_edges_naive = edges_packed.grad.cpu()
 
         # Compare
         self.assertClose(grad_points_naive.cpu(), grad_points_cuda.cpu(), atol=1e-7)
-        self.assertClose(grad_edges_naive.cpu(), grad_edges_cuda.cpu(), atol=5e-7)
+        self.assertClose(grad_edges_naive, grad_edges_cuda.cpu(), atol=5e-7)
+        self.assertClose(grad_points_naive.cpu(), grad_points_cpu, atol=1e-7)
+        self.assertClose(grad_edges_naive, grad_edges_cpu, atol=5e-7)
 
     def test_edge_point_distance(self):
         """
@@ -361,6 +388,20 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
             points_packed, edges_packed, idx_cuda, grad_dists
         )
 
+        # Cpu Implementation: forward
+        dists_cpu, idx_cpu = _C.edge_point_dist_forward(
+            points_packed.cpu(),
+            points_first_idx.cpu(),
+            edges_packed.cpu(),
+            edges_first_idx.cpu(),
+            max_e,
+        )
+
+        # Cpu Implementation: backward
+        grad_points_cpu, grad_edges_cpu = _C.edge_point_dist_backward(
+            points_packed.cpu(), edges_packed.cpu(), idx_cpu, grad_dists.cpu()
+        )
+
         # Naive Implementation: forward
         edges_list = packed_to_list(edges_packed, meshes.num_edges_per_mesh().tolist())
         dists_naive = []
@@ -395,15 +436,18 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
 
         # Compare
         self.assertClose(dists_naive.cpu(), dists_cuda.cpu())
+        self.assertClose(dists_naive.cpu(), dists_cpu)
 
         # Naive Implementation: backward
         dists_naive.backward(grad_dists)
         grad_points_naive = torch.cat([cloud.grad for cloud in pcls.points_list()])
-        grad_edges_naive = edges_packed.grad
+        grad_edges_naive = edges_packed.grad.cpu()
 
         # Compare
         self.assertClose(grad_points_naive.cpu(), grad_points_cuda.cpu(), atol=1e-7)
-        self.assertClose(grad_edges_naive.cpu(), grad_edges_cuda.cpu(), atol=5e-7)
+        self.assertClose(grad_edges_naive, grad_edges_cuda.cpu(), atol=5e-7)
+        self.assertClose(grad_points_naive.cpu(), grad_points_cpu, atol=1e-7)
+        self.assertClose(grad_edges_naive, grad_edges_cpu, atol=5e-7)
 
     def test_point_mesh_edge_distance(self):
         """
@@ -483,6 +527,8 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
         device = get_random_cuda_device()
         points = torch.rand((P, 3), dtype=torch.float32, device=device)
         tris = torch.rand((T, 3, 3), dtype=torch.float32, device=device)
+        points_cpu = points.clone().cpu()
+        tris_cpu = tris.clone().cpu()
 
         points.requires_grad = True
         tris.requires_grad = True
@@ -502,23 +548,30 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
 
         # Naive Backward
         dists_naive.backward(grad_dists)
-        grad_points_naive = points.grad
-        grad_tris_naive = tris.grad
+        grad_points_naive = points.grad.cpu()
+        grad_tris_naive = tris.grad.cpu()
 
         # Cuda Forward Implementation
         dists_cuda = _C.point_face_array_dist_forward(points, tris)
+        dists_cpu = _C.point_face_array_dist_forward(points_cpu, tris_cpu)
 
         # Compare
         self.assertClose(dists_naive.cpu(), dists_cuda.cpu())
+        self.assertClose(dists_naive.cpu(), dists_cpu)
 
         # CUDA Backward Implementation
         grad_points_cuda, grad_tris_cuda = _C.point_face_array_dist_backward(
             points, tris, grad_dists
         )
+        grad_points_cpu, grad_tris_cpu = _C.point_face_array_dist_backward(
+            points_cpu, tris_cpu, grad_dists.cpu()
+        )
 
         # Compare
-        self.assertClose(grad_points_naive.cpu(), grad_points_cuda.cpu())
-        self.assertClose(grad_tris_naive.cpu(), grad_tris_cuda.cpu(), atol=5e-6)
+        self.assertClose(grad_points_naive, grad_points_cuda.cpu())
+        self.assertClose(grad_tris_naive, grad_tris_cuda.cpu(), atol=5e-6)
+        self.assertClose(grad_points_naive, grad_points_cpu)
+        self.assertClose(grad_tris_naive, grad_tris_cpu, atol=5e-6)
 
     def test_point_face_distance(self):
         """
@@ -559,6 +612,21 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
             points_packed, faces_packed, idx_cuda, grad_dists
         )
 
+        # Cpu Implementation: forward
+        dists_cpu, idx_cpu = _C.point_face_dist_forward(
+            points_packed.cpu(),
+            points_first_idx.cpu(),
+            faces_packed.cpu(),
+            faces_first_idx.cpu(),
+            max_p,
+        )
+
+        # Cpu Implementation: backward
+        # Note that using idx_cpu doesn't pass - there seems to be a problem with tied results.
+        grad_points_cpu, grad_faces_cpu = _C.point_face_dist_backward(
+            points_packed.cpu(), faces_packed.cpu(), idx_cuda.cpu(), grad_dists.cpu()
+        )
+
         # Naive Implementation: forward
         faces_list = packed_to_list(faces_packed, meshes.num_faces_per_mesh().tolist())
         dists_naive = []
@@ -593,15 +661,18 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
 
         # Compare
         self.assertClose(dists_naive.cpu(), dists_cuda.cpu())
+        self.assertClose(dists_naive.cpu(), dists_cpu)
 
         #  Naive Implementation: backward
         dists_naive.backward(grad_dists)
         grad_points_naive = torch.cat([cloud.grad for cloud in pcls.points_list()])
-        grad_faces_naive = faces_packed.grad
+        grad_faces_naive = faces_packed.grad.cpu()
 
         # Compare
         self.assertClose(grad_points_naive.cpu(), grad_points_cuda.cpu(), atol=1e-7)
-        self.assertClose(grad_faces_naive.cpu(), grad_faces_cuda.cpu(), atol=5e-7)
+        self.assertClose(grad_faces_naive, grad_faces_cuda.cpu(), atol=5e-7)
+        self.assertClose(grad_points_naive.cpu(), grad_points_cpu, atol=1e-7)
+        self.assertClose(grad_faces_naive, grad_faces_cpu, atol=5e-7)
 
     def test_face_point_distance(self):
         """
@@ -642,6 +713,20 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
             points_packed, faces_packed, idx_cuda, grad_dists
         )
 
+        # Cpu Implementation: forward
+        dists_cpu, idx_cpu = _C.face_point_dist_forward(
+            points_packed.cpu(),
+            points_first_idx.cpu(),
+            faces_packed.cpu(),
+            faces_first_idx.cpu(),
+            max_f,
+        )
+
+        # Cpu Implementation: backward
+        grad_points_cpu, grad_faces_cpu = _C.face_point_dist_backward(
+            points_packed.cpu(), faces_packed.cpu(), idx_cpu, grad_dists.cpu()
+        )
+
         # Naive Implementation: forward
         faces_list = packed_to_list(faces_packed, meshes.num_faces_per_mesh().tolist())
         dists_naive = []
@@ -676,6 +761,7 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
 
         # Compare
         self.assertClose(dists_naive.cpu(), dists_cuda.cpu())
+        self.assertClose(dists_naive.cpu(), dists_cpu)
 
         # Naive Implementation: backward
         dists_naive.backward(grad_dists)
@@ -685,6 +771,8 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
         # Compare
         self.assertClose(grad_points_naive.cpu(), grad_points_cuda.cpu(), atol=1e-7)
         self.assertClose(grad_faces_naive.cpu(), grad_faces_cuda.cpu(), atol=5e-7)
+        self.assertClose(grad_points_naive.cpu(), grad_points_cpu, atol=1e-7)
+        self.assertClose(grad_faces_naive.cpu(), grad_faces_cpu, atol=5e-7)
 
     def test_point_mesh_face_distance(self):
         """