From 3ff6c5ab85c637713b2cffddef2b3f4710a8d1cc Mon Sep 17 00:00:00 2001
From: Melvin He <melvinhe@meta.com>
Date: Tue, 1 Jul 2025 09:14:38 -0700
Subject: [PATCH] Error instead of crash for tensors on exotic devices

Summary:
Adds CHECK_CPU macros that checks if a tensor is on the CPU device throughout csrc directories up to `marching_cubes`. Directories updated include those in `gather_scatter`, `interp_face_attrs`, `iou_box3d`, `knn`, and `marching_cubes`.

Note that this is the second part of a larger change, and to keep diffs better organized, subsequent diffs will update the remaining directories.

Reviewed By: bottler

Differential Revision: D77558550

fbshipit-source-id: 762a0fe88548dc8d0901b198a11c40d0c36e173f
---
 pytorch3d/csrc/gather_scatter/gather_scatter.h       | 2 ++
 pytorch3d/csrc/interp_face_attrs/interp_face_attrs.h | 5 +++++
 pytorch3d/csrc/iou_box3d/iou_box3d.h                 | 2 ++
 pytorch3d/csrc/knn/knn.h                             | 4 ++++
 pytorch3d/csrc/marching_cubes/marching_cubes.h       | 1 +
 5 files changed, 14 insertions(+)

diff --git a/pytorch3d/csrc/gather_scatter/gather_scatter.h b/pytorch3d/csrc/gather_scatter/gather_scatter.h
index 9ab9574f..21d5f7b1 100644
--- a/pytorch3d/csrc/gather_scatter/gather_scatter.h
+++ b/pytorch3d/csrc/gather_scatter/gather_scatter.h
@@ -53,5 +53,7 @@ at::Tensor GatherScatter(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(input);
+  CHECK_CPU(edges);
   return GatherScatterCpu(input, edges, directed, backward);
 }
diff --git a/pytorch3d/csrc/interp_face_attrs/interp_face_attrs.h b/pytorch3d/csrc/interp_face_attrs/interp_face_attrs.h
index 5ba14462..42f93230 100644
--- a/pytorch3d/csrc/interp_face_attrs/interp_face_attrs.h
+++ b/pytorch3d/csrc/interp_face_attrs/interp_face_attrs.h
@@ -57,6 +57,8 @@ at::Tensor InterpFaceAttrsForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(face_attrs);
+  CHECK_CPU(barycentric_coords);
   return InterpFaceAttrsForwardCpu(pix_to_face, barycentric_coords, face_attrs);
 }
 
@@ -106,6 +108,9 @@ std::tuple<at::Tensor, at::Tensor> InterpFaceAttrsBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(face_attrs);
+  CHECK_CPU(barycentric_coords);
+  CHECK_CPU(grad_pix_attrs);
   return InterpFaceAttrsBackwardCpu(
       pix_to_face, barycentric_coords, face_attrs, grad_pix_attrs);
 }
diff --git a/pytorch3d/csrc/iou_box3d/iou_box3d.h b/pytorch3d/csrc/iou_box3d/iou_box3d.h
index 84f752b0..d5034b3d 100644
--- a/pytorch3d/csrc/iou_box3d/iou_box3d.h
+++ b/pytorch3d/csrc/iou_box3d/iou_box3d.h
@@ -44,5 +44,7 @@ inline std::tuple<at::Tensor, at::Tensor> IoUBox3D(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(boxes1);
+  CHECK_CPU(boxes2);
   return IoUBox3DCpu(boxes1.contiguous(), boxes2.contiguous());
 }
diff --git a/pytorch3d/csrc/knn/knn.h b/pytorch3d/csrc/knn/knn.h
index 7fc8d488..63f204cd 100644
--- a/pytorch3d/csrc/knn/knn.h
+++ b/pytorch3d/csrc/knn/knn.h
@@ -74,6 +74,8 @@ std::tuple<at::Tensor, at::Tensor> KNearestNeighborIdx(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(p1);
+  CHECK_CPU(p2);
   return KNearestNeighborIdxCpu(p1, p2, lengths1, lengths2, norm, K);
 }
 
@@ -140,6 +142,8 @@ std::tuple<at::Tensor, at::Tensor> KNearestNeighborBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(p1);
+  CHECK_CPU(p2);
   return KNearestNeighborBackwardCpu(
       p1, p2, lengths1, lengths2, idxs, norm, grad_dists);
 }
diff --git a/pytorch3d/csrc/marching_cubes/marching_cubes.h b/pytorch3d/csrc/marching_cubes/marching_cubes.h
index 51c660b1..7a425c07 100644
--- a/pytorch3d/csrc/marching_cubes/marching_cubes.h
+++ b/pytorch3d/csrc/marching_cubes/marching_cubes.h
@@ -58,5 +58,6 @@ inline std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubes(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(vol);
   return MarchingCubesCpu(vol.contiguous(), isolevel);
 }