increment_version for inplace ops

Summary: For safety checks, make inplace forward operations in cuda and c++ call increment_version. Reviewed By: davidsonic Differential Revision: D44302504 fbshipit-source-id: 6ff62251e352d6778cb54399e2e11459e16e77ba
facebookresearch · Mar 23, 2023 · 7d8b029 · 7d8b029
1 parent 9437768
commit 7d8b029
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 0 deletions.
diff --git a/pytorch3d/csrc/points_to_volumes/points_to_volumes.h b/pytorch3d/csrc/points_to_volumes/points_to_volumes.h
@@ -7,6 +7,7 @@
  */
 
 #pragma once
+#include <torch/csrc/autograd/VariableTypeUtils.h>
 #include <torch/extension.h>
 #include <cstdio>
 #include <tuple>
@@ -96,6 +97,8 @@ inline void PointsToVolumesForward(
         point_weight,
         align_corners,
         splat);
+    torch::autograd::increment_version(volume_features);
+    torch::autograd::increment_version(volume_densities);
     return;
 #else
     AT_ERROR("Not compiled with GPU support.");

diff --git a/pytorch3d/csrc/points_to_volumes/points_to_volumes_cpu.cpp b/pytorch3d/csrc/points_to_volumes/points_to_volumes_cpu.cpp
@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <torch/csrc/autograd/VariableTypeUtils.h>
 #include <torch/extension.h>
 #include <algorithm>
 #include <cmath>
@@ -148,6 +149,8 @@ void PointsToVolumesForwardCpu(
       }
     }
   }
+  torch::autograd::increment_version(volume_features);
+  torch::autograd::increment_version(volume_densities);
 }
 
 // With nearest, the only smooth dependence is that volume features

diff --git a/pytorch3d/csrc/sample_pdf/sample_pdf.h b/pytorch3d/csrc/sample_pdf/sample_pdf.h
@@ -7,6 +7,7 @@
  */
 
 #pragma once
+#include <torch/csrc/autograd/VariableTypeUtils.h>
 #include <torch/extension.h>
 #include <cstdio>
 #include <tuple>
@@ -63,6 +64,7 @@ inline void SamplePdf(
 #ifdef WITH_CUDA
     CHECK_CUDA(weights);
     CHECK_CONTIGUOUS_CUDA(outputs);
+    torch::autograd::increment_version(outputs);
     SamplePdfCuda(bins, weights, outputs, eps);
     return;
 #else

diff --git a/pytorch3d/csrc/sample_pdf/sample_pdf_cpu.cpp b/pytorch3d/csrc/sample_pdf/sample_pdf_cpu.cpp
@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <torch/csrc/autograd/VariableTypeUtils.h>
 #include <torch/extension.h>
 #include <algorithm>
 #include <thread>
@@ -137,4 +138,5 @@ void SamplePdfCpu(
   for (auto&& thread : threads) {
     thread.join();
   }
+  torch::autograd::increment_version(outputs);
 }