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
This commit is contained in:
Jeremy Reizenstein 2023-03-23 11:48:36 -07:00 committed by Facebook GitHub Bot
parent 9437768687
commit 7d8b029aae
4 changed files with 10 additions and 0 deletions

View File

@ -7,6 +7,7 @@
*/ */
#pragma once #pragma once
#include <torch/csrc/autograd/VariableTypeUtils.h>
#include <torch/extension.h> #include <torch/extension.h>
#include <cstdio> #include <cstdio>
#include <tuple> #include <tuple>
@ -96,6 +97,8 @@ inline void PointsToVolumesForward(
point_weight, point_weight,
align_corners, align_corners,
splat); splat);
torch::autograd::increment_version(volume_features);
torch::autograd::increment_version(volume_densities);
return; return;
#else #else
AT_ERROR("Not compiled with GPU support."); AT_ERROR("Not compiled with GPU support.");

View File

@ -6,6 +6,7 @@
* LICENSE file in the root directory of this source tree. * LICENSE file in the root directory of this source tree.
*/ */
#include <torch/csrc/autograd/VariableTypeUtils.h>
#include <torch/extension.h> #include <torch/extension.h>
#include <algorithm> #include <algorithm>
#include <cmath> #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 // With nearest, the only smooth dependence is that volume features

View File

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

View File

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