From 177eec63785ac14aaa674b3b132dcdbf98a7afad Mon Sep 17 00:00:00 2001 From: Melvin He Date: Mon, 30 Jun 2025 12:27:38 -0700 Subject: [PATCH] Error instead of crash for tensors on exotic devices (#1986) Summary: Pull Request resolved: https://github.com/facebookresearch/pytorch3d/pull/1986 Adds device checks to prevent crashes on unsupported devices in PyTorch3D. Updates the `pytorch3d_cutils.h` file to include new macro CHECK_CPU that checks if a tensor is on the CPU device. This macro is then used in the directories from `ball_query` to `face_area_normals` to ensure that tensors are not on unsupported devices like MPS. Note that this is the first part of a larger change, and to keep diffs better organized, subsequent diffs will update the remaining directories. Reviewed By: bottler Differential Revision: D77473296 fbshipit-source-id: 13dc84620dee667bddebad1dade2d2cb5a59c737 --- pytorch3d/csrc/ball_query/ball_query.h | 2 ++ pytorch3d/csrc/blending/sigmoid_alpha_blend.h | 5 +++++ pytorch3d/csrc/compositing/alpha_composite.h | 8 ++++++++ pytorch3d/csrc/compositing/norm_weighted_sum.h | 9 +++++++++ pytorch3d/csrc/compositing/weighted_sum.h | 8 ++++++++ pytorch3d/csrc/face_areas_normals/face_areas_normals.h | 6 ++++++ pytorch3d/csrc/utils/pytorch3d_cutils.h | 4 ++++ 7 files changed, 42 insertions(+) diff --git a/pytorch3d/csrc/ball_query/ball_query.h b/pytorch3d/csrc/ball_query/ball_query.h index 059cad8b..eb8f54da 100644 --- a/pytorch3d/csrc/ball_query/ball_query.h +++ b/pytorch3d/csrc/ball_query/ball_query.h @@ -81,6 +81,8 @@ inline std::tuple BallQuery( AT_ERROR("Not compiled with GPU support."); #endif } + CHECK_CPU(p1); + CHECK_CPU(p2); return BallQueryCpu( p1.contiguous(), p2.contiguous(), diff --git a/pytorch3d/csrc/blending/sigmoid_alpha_blend.h b/pytorch3d/csrc/blending/sigmoid_alpha_blend.h index d424c769..8b023b61 100644 --- a/pytorch3d/csrc/blending/sigmoid_alpha_blend.h +++ b/pytorch3d/csrc/blending/sigmoid_alpha_blend.h @@ -98,6 +98,11 @@ at::Tensor SigmoidAlphaBlendBackward( AT_ERROR("Not compiled with GPU support."); #endif } + CHECK_CPU(distances); + CHECK_CPU(pix_to_face); + CHECK_CPU(alphas); + CHECK_CPU(grad_alphas); + return SigmoidAlphaBlendBackwardCpu( grad_alphas, alphas, distances, pix_to_face, sigma); } diff --git a/pytorch3d/csrc/compositing/alpha_composite.h b/pytorch3d/csrc/compositing/alpha_composite.h index 44ba1bf0..a9ec7b43 100644 --- a/pytorch3d/csrc/compositing/alpha_composite.h +++ b/pytorch3d/csrc/compositing/alpha_composite.h @@ -74,6 +74,9 @@ torch::Tensor alphaCompositeForward( AT_ERROR("Not compiled with GPU support"); #endif } else { + CHECK_CPU(features); + CHECK_CPU(alphas); + CHECK_CPU(points_idx); return alphaCompositeCpuForward(features, alphas, points_idx); } } @@ -101,6 +104,11 @@ std::tuple alphaCompositeBackward( AT_ERROR("Not compiled with GPU support"); #endif } else { + CHECK_CPU(grad_outputs); + CHECK_CPU(features); + CHECK_CPU(alphas); + CHECK_CPU(points_idx); + return alphaCompositeCpuBackward( grad_outputs, features, alphas, points_idx); } diff --git a/pytorch3d/csrc/compositing/norm_weighted_sum.h b/pytorch3d/csrc/compositing/norm_weighted_sum.h index 5d0a5f5b..5f72ee2d 100644 --- a/pytorch3d/csrc/compositing/norm_weighted_sum.h +++ b/pytorch3d/csrc/compositing/norm_weighted_sum.h @@ -73,6 +73,10 @@ torch::Tensor weightedSumNormForward( AT_ERROR("Not compiled with GPU support"); #endif } else { + CHECK_CPU(features); + CHECK_CPU(alphas); + CHECK_CPU(points_idx); + return weightedSumNormCpuForward(features, alphas, points_idx); } } @@ -100,6 +104,11 @@ std::tuple weightedSumNormBackward( AT_ERROR("Not compiled with GPU support"); #endif } else { + CHECK_CPU(grad_outputs); + CHECK_CPU(features); + CHECK_CPU(alphas); + CHECK_CPU(points_idx); + return weightedSumNormCpuBackward( grad_outputs, features, alphas, points_idx); } diff --git a/pytorch3d/csrc/compositing/weighted_sum.h b/pytorch3d/csrc/compositing/weighted_sum.h index 0be6e498..cdc3fdf5 100644 --- a/pytorch3d/csrc/compositing/weighted_sum.h +++ b/pytorch3d/csrc/compositing/weighted_sum.h @@ -72,6 +72,9 @@ torch::Tensor weightedSumForward( AT_ERROR("Not compiled with GPU support"); #endif } else { + CHECK_CPU(features); + CHECK_CPU(alphas); + CHECK_CPU(points_idx); return weightedSumCpuForward(features, alphas, points_idx); } } @@ -98,6 +101,11 @@ std::tuple weightedSumBackward( AT_ERROR("Not compiled with GPU support"); #endif } else { + CHECK_CPU(grad_outputs); + CHECK_CPU(features); + CHECK_CPU(alphas); + CHECK_CPU(points_idx); + return weightedSumCpuBackward(grad_outputs, features, alphas, points_idx); } } diff --git a/pytorch3d/csrc/face_areas_normals/face_areas_normals.h b/pytorch3d/csrc/face_areas_normals/face_areas_normals.h index 6df37c12..201abd28 100644 --- a/pytorch3d/csrc/face_areas_normals/face_areas_normals.h +++ b/pytorch3d/csrc/face_areas_normals/face_areas_normals.h @@ -60,6 +60,8 @@ std::tuple FaceAreasNormalsForward( AT_ERROR("Not compiled with GPU support."); #endif } + CHECK_CPU(verts); + CHECK_CPU(faces); return FaceAreasNormalsForwardCpu(verts, faces); } @@ -80,5 +82,9 @@ at::Tensor FaceAreasNormalsBackward( AT_ERROR("Not compiled with GPU support."); #endif } + CHECK_CPU(grad_areas); + CHECK_CPU(grad_normals); + CHECK_CPU(verts); + CHECK_CPU(faces); return FaceAreasNormalsBackwardCpu(grad_areas, grad_normals, verts, faces); } diff --git a/pytorch3d/csrc/utils/pytorch3d_cutils.h b/pytorch3d/csrc/utils/pytorch3d_cutils.h index 48d04546..c1856eca 100644 --- a/pytorch3d/csrc/utils/pytorch3d_cutils.h +++ b/pytorch3d/csrc/utils/pytorch3d_cutils.h @@ -15,3 +15,7 @@ #define CHECK_CONTIGUOUS_CUDA(x) \ CHECK_CUDA(x); \ CHECK_CONTIGUOUS(x) +#define CHECK_CPU(x) \ + TORCH_CHECK( \ + x.device().type() == torch::kCPU, \ + "Cannot use CPU implementation: " #x " not on CPU.")