[UPDATE] remove grad_face_vertices

install.sh

@@ -1,2 +1,6 @@
+if [ -d "build" ];then
+    bash clean.sh
+fi
+
 python setup.py develop
 python -c "import torchsdf; print(torchsdf.__version__)"

test.sh

@@ -0,0 +1 @@
+python tests/value.py

tests/value.py

@@ -7,7 +7,7 @@
 os.environ["CUDA_VISIBLE_DIVICES"] = "1"
 device = "cuda"
 # Ns
-num_sample = 10000
+num_sample = 1000000
 
 all_pass = True
 

torchsdf/csrc/unbatched_triangle_distance.cpp

@@ -34,8 +34,7 @@ void unbatched_triangle_distance_backward_cuda_impl(
     at::Tensor face_vertices,
     at::Tensor face_idx,
     at::Tensor dist_type,
-    at::Tensor grad_points,
-    at::Tensor grad_face_vertices);
+    at::Tensor grad_points);
 
 #endif  // WITH_CUDA
 
@@ -77,22 +76,19 @@ void unbatched_triangle_distance_backward_cuda(
     at::Tensor face_vertices,
     at::Tensor face_idx,
     at::Tensor dist_type,
-    at::Tensor grad_points,
-    at::Tensor grad_face_vertices) {
+    at::Tensor grad_points) {
   CHECK_CUDA(grad_dist);
   CHECK_CUDA(points);
   CHECK_CUDA(face_vertices);
   CHECK_CUDA(face_idx);
   CHECK_CUDA(dist_type);
   CHECK_CUDA(grad_points);
-  CHECK_CUDA(grad_face_vertices);
   CHECK_CONTIGUOUS(grad_dist);
   CHECK_CONTIGUOUS(points);
   CHECK_CONTIGUOUS(face_vertices);
   CHECK_CONTIGUOUS(face_idx);
   CHECK_CONTIGUOUS(dist_type);
   CHECK_CONTIGUOUS(grad_points);
-  CHECK_CONTIGUOUS(grad_face_vertices);
 
   const int num_points = points.size(0);
   const int num_faces = face_vertices.size(0);
@@ -102,12 +98,11 @@ void unbatched_triangle_distance_backward_cuda(
   CHECK_SIZES(face_idx, num_points);
   CHECK_SIZES(dist_type, num_points);
   CHECK_SIZES(grad_points, num_points, 3);
-  CHECK_SIZES(grad_face_vertices, num_faces, 3, 3);
 
 #if WITH_CUDA
   unbatched_triangle_distance_backward_cuda_impl(
       grad_dist, points, face_vertices, face_idx, dist_type,
-      grad_points, grad_face_vertices);
+      grad_points);
 #else
   AT_ERROR("unbatched_triangle_distance_backward not built with CUDA");
 #endif

torchsdf/csrc/unbatched_triangle_distance.h

@@ -33,8 +33,7 @@ void unbatched_triangle_distance_backward_cuda(
     at::Tensor face_vertices,
     at::Tensor face_idx,
     at::Tensor dist_type,
-    at::Tensor grad_points,
-    at::Tensor grad_face_vertices);
+    at::Tensor grad_points);
 
 }  // namespace kaolin
 

torchsdf/csrc/unbatched_triangle_distance_cuda.cu

@@ -176,8 +176,6 @@ template<typename scalar_t, typename vector_t>
 __device__ __forceinline__ void compute_edge_backward(
     vector_t vab,
     vector_t pb,
-    scalar_t* grad_input_va,
-    scalar_t* grad_input_vb,
     scalar_t* grad_input_p,
     int64_t index,
     scalar_t grad) {
@@ -221,14 +219,6 @@ __device__ __forceinline__ void compute_edge_backward(
   grad_input_p[index * 3] = p_bar.x;
   grad_input_p[index * 3 + 1] = p_bar.y;
   grad_input_p[index * 3 + 2] = p_bar.z;
-
-  atomicAdd(&(grad_input_va[0]), va_bar.x);
-  atomicAdd(&(grad_input_va[1]), va_bar.y);
-  atomicAdd(&(grad_input_va[2]), va_bar.z);
-
-  atomicAdd(&(grad_input_vb[0]), vb_bar.x);
-  atomicAdd(&(grad_input_vb[1]), vb_bar.y);
-  atomicAdd(&(grad_input_vb[2]), vb_bar.z);
 }
 
 
@@ -324,8 +314,7 @@ __global__ void unbatched_triangle_distance_backward_cuda_kernel(
     int* dist_type,
     int num_points,
     int num_faces,
-    scalar_t* grad_points,
-    scalar_t* grad_face_vertices) {
+    scalar_t* grad_points) {
   for (int point_id = threadIdx.x + blockIdx.x * blockDim.x; point_id < num_points;
        point_id += blockDim.x * gridDim.x) {
     int type = dist_type[point_id];
@@ -361,53 +350,29 @@ __global__ void unbatched_triangle_distance_backward_cuda_kernel(
       grad_points[point_id * 3 + 1] = grad_point_vec.y;
       grad_points[point_id * 3 + 2] = grad_point_vec.z;
       vector_t tmp = grad_e31 + grad_e21 - grad_point_vec;
-      atomicAdd(&(grad_face_vertices[face_id * 9]), tmp.x);
-      atomicAdd(&(grad_face_vertices[face_id * 9 + 1]), tmp.y);
-      atomicAdd(&(grad_face_vertices[face_id * 9 + 2]), tmp.z);
-      atomicAdd(&(grad_face_vertices[face_id * 9 + 3]), -grad_e21.x);
-      atomicAdd(&(grad_face_vertices[face_id * 9 + 4]), -grad_e21.y);
-      atomicAdd(&(grad_face_vertices[face_id * 9 + 5]), -grad_e21.z);
-      atomicAdd(&(grad_face_vertices[face_id * 9 + 6]), -grad_e31.x);
-      atomicAdd(&(grad_face_vertices[face_id * 9 + 7]), -grad_e31.y);
-      atomicAdd(&(grad_face_vertices[face_id * 9 + 8]), -grad_e31.z);
     } else if (type == 1) {  // distance to v1
       vector_t grad_dist_vec = (p - v1) * grad_out;
-      atomicAdd(&(grad_face_vertices[face_id * 9]), -grad_dist_vec.x);
-      atomicAdd(&(grad_face_vertices[face_id * 9 + 1]), -grad_dist_vec.y);
-      atomicAdd(&(grad_face_vertices[face_id * 9 + 2]), -grad_dist_vec.z);
       grad_points[point_id * 3] = grad_dist_vec.x;
       grad_points[point_id * 3 + 1] = grad_dist_vec.y;
       grad_points[point_id * 3 + 2] = grad_dist_vec.z;
     } else if (type == 2) {  // distance to v2
       vector_t grad_dist_vec = (p - v2) * grad_out;
-      atomicAdd(&(grad_face_vertices[face_id * 9 + 3]), -grad_dist_vec.x);
-      atomicAdd(&(grad_face_vertices[face_id * 9 + 4]), -grad_dist_vec.y);
-      atomicAdd(&(grad_face_vertices[face_id * 9 + 5]), -grad_dist_vec.z);
       grad_points[point_id * 3] = grad_dist_vec.x;
       grad_points[point_id * 3 + 1] = grad_dist_vec.y;
       grad_points[point_id * 3 + 2] = grad_dist_vec.z;
     } else if (type == 3) {  // distance to v3
       vector_t grad_dist_vec = (p - v3) * grad_out;
-      atomicAdd(&(grad_face_vertices[face_id * 9 + 6]), -grad_dist_vec.x);
-      atomicAdd(&(grad_face_vertices[face_id * 9 + 7]), -grad_dist_vec.y);
-      atomicAdd(&(grad_face_vertices[face_id * 9 + 8]), -grad_dist_vec.z);
       grad_points[point_id * 3] = grad_dist_vec.x;
       grad_points[point_id * 3 + 1] = grad_dist_vec.y;
       grad_points[point_id * 3 + 2] = grad_dist_vec.z;
     } else if (type == 4) {  // distance to e12
       compute_edge_backward(e12, p - v1,
-                            &(grad_face_vertices[face_id * 9 + 3]),
-                            &(grad_face_vertices[face_id * 9]),
                             grad_points, point_id, grad_out);
     } else if (type == 5) {  // distance to e23
       compute_edge_backward(e23, p - v2,
-                            &(grad_face_vertices[face_id * 9 + 6]),
-                            &(grad_face_vertices[face_id * 9 + 3]),
                             grad_points, point_id, grad_out);
     } else {  // distance to e31
       compute_edge_backward(e31, p - v3,
-                            &(grad_face_vertices[face_id * 9]),
-                            &(grad_face_vertices[face_id * 9 + 6]),
                             grad_points, point_id, grad_out);
     }
   }
@@ -447,8 +412,7 @@ void unbatched_triangle_distance_backward_cuda_impl(
     at::Tensor face_vertices,
     at::Tensor face_idx,
     at::Tensor dist_type,
-    at::Tensor grad_points,
-    at::Tensor grad_face_vertices) {
+    at::Tensor grad_points) {
 
   DISPATCH_INPUT_TYPES(points.scalar_type(), scalar_t,
                        "unbatched_triangle_distance_backward_cuda", [&] {
@@ -466,8 +430,7 @@ void unbatched_triangle_distance_backward_cuda_impl(
         dist_type.data_ptr<int32_t>(),
         points.size(0),
         face_vertices.size(0),
-        grad_points.data_ptr<scalar_t>(),
-        grad_face_vertices.data_ptr<scalar_t>());
+        grad_points.data_ptr<scalar_t>());
     CUDA_CHECK(cudaGetLastError());
   });
 }

torchsdf/sdf.py

@@ -58,8 +58,8 @@ def backward(ctx, grad_dist, grad_face_idx, grad_dist_type):
         points, face_vertices, face_idx, dist_type = ctx.saved_tensors
         grad_dist = grad_dist.contiguous()
         grad_points = torch.zeros_like(points)
-        grad_face_vertices = torch.zeros_like(face_vertices)
+        grad_face_vertices = None
         _C.unbatched_triangle_distance_backward_cuda(
             grad_dist, points, face_vertices, face_idx, dist_type,
-            grad_points, grad_face_vertices)
+            grad_points)
         return grad_points, grad_face_vertices