From 7d8b029aaed5f247e3898f4b51e6a79ca52bfc54 Mon Sep 17 00:00:00 2001
From: Jeremy Reizenstein <reizenstein@meta.com>
Date: Thu, 23 Mar 2023 11:48:36 -0700
Subject: [PATCH] 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
---
 pytorch3d/csrc/points_to_volumes/points_to_volumes.h       | 3 +++
 pytorch3d/csrc/points_to_volumes/points_to_volumes_cpu.cpp | 3 +++
 pytorch3d/csrc/sample_pdf/sample_pdf.h                     | 2 ++
 pytorch3d/csrc/sample_pdf/sample_pdf_cpu.cpp               | 2 ++
 4 files changed, 10 insertions(+)

diff --git a/pytorch3d/csrc/points_to_volumes/points_to_volumes.h b/pytorch3d/csrc/points_to_volumes/points_to_volumes.h
index 9b077066..4c5eba3c 100644
--- 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
index ac7a3d98..900ea097 100644
--- 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
index 66775c4f..899117df 100644
--- 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
index b297ae9e..272197c6 100644
--- 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);
 }