Fixes bus error hard crashes on Apple Silicon MPS devices

Summary:
Fixes hard crashes (bus errors) when using MPS device (Apple Silicon) by implementing CPU checks throughout files in csrc subdirectories to check if on same mesh on a CPU device.

Note that this is the fourth and ultimate part of a larger change through multiple files & directories.

Reviewed By: bottler

Differential Revision: D77698176

fbshipit-source-id: 5bc9e3c5cea61afd486aed7396f390d92775ec6d
This commit is contained in:
Melvin He 2025-07-03 12:34:37 -07:00 committed by Facebook GitHub Bot
parent c5ea8fa49e
commit 3aee2a6005
4 changed files with 33 additions and 0 deletions

View File

@ -138,6 +138,9 @@ RasterizeMeshesNaive(
AT_ERROR("Not compiled with GPU support"); AT_ERROR("Not compiled with GPU support");
#endif #endif
} else { } else {
CHECK_CPU(face_verts);
CHECK_CPU(mesh_to_face_first_idx);
CHECK_CPU(num_faces_per_mesh);
return RasterizeMeshesNaiveCpu( return RasterizeMeshesNaiveCpu(
face_verts, face_verts,
mesh_to_face_first_idx, mesh_to_face_first_idx,
@ -232,6 +235,11 @@ torch::Tensor RasterizeMeshesBackward(
AT_ERROR("Not compiled with GPU support"); AT_ERROR("Not compiled with GPU support");
#endif #endif
} else { } else {
CHECK_CPU(face_verts);
CHECK_CPU(pix_to_face);
CHECK_CPU(grad_zbuf);
CHECK_CPU(grad_bary);
CHECK_CPU(grad_dists);
return RasterizeMeshesBackwardCpu( return RasterizeMeshesBackwardCpu(
face_verts, face_verts,
pix_to_face, pix_to_face,
@ -306,6 +314,9 @@ torch::Tensor RasterizeMeshesCoarse(
AT_ERROR("Not compiled with GPU support"); AT_ERROR("Not compiled with GPU support");
#endif #endif
} else { } else {
CHECK_CPU(face_verts);
CHECK_CPU(mesh_to_face_first_idx);
CHECK_CPU(num_faces_per_mesh);
return RasterizeMeshesCoarseCpu( return RasterizeMeshesCoarseCpu(
face_verts, face_verts,
mesh_to_face_first_idx, mesh_to_face_first_idx,
@ -423,6 +434,8 @@ RasterizeMeshesFine(
AT_ERROR("Not compiled with GPU support"); AT_ERROR("Not compiled with GPU support");
#endif #endif
} else { } else {
CHECK_CPU(face_verts);
CHECK_CPU(bin_faces);
AT_ERROR("NOT IMPLEMENTED"); AT_ERROR("NOT IMPLEMENTED");
} }
} }

View File

@ -91,6 +91,10 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> RasterizePointsNaive(
AT_ERROR("Not compiled with GPU support"); AT_ERROR("Not compiled with GPU support");
#endif #endif
} else { } else {
CHECK_CPU(points);
CHECK_CPU(cloud_to_packed_first_idx);
CHECK_CPU(num_points_per_cloud);
CHECK_CPU(radius);
return RasterizePointsNaiveCpu( return RasterizePointsNaiveCpu(
points, points,
cloud_to_packed_first_idx, cloud_to_packed_first_idx,
@ -166,6 +170,10 @@ torch::Tensor RasterizePointsCoarse(
AT_ERROR("Not compiled with GPU support"); AT_ERROR("Not compiled with GPU support");
#endif #endif
} else { } else {
CHECK_CPU(points);
CHECK_CPU(cloud_to_packed_first_idx);
CHECK_CPU(num_points_per_cloud);
CHECK_CPU(radius);
return RasterizePointsCoarseCpu( return RasterizePointsCoarseCpu(
points, points,
cloud_to_packed_first_idx, cloud_to_packed_first_idx,
@ -232,6 +240,8 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> RasterizePointsFine(
AT_ERROR("Not compiled with GPU support"); AT_ERROR("Not compiled with GPU support");
#endif #endif
} else { } else {
CHECK_CPU(points);
CHECK_CPU(bin_points);
AT_ERROR("NOT IMPLEMENTED"); AT_ERROR("NOT IMPLEMENTED");
} }
} }
@ -284,6 +294,10 @@ torch::Tensor RasterizePointsBackward(
AT_ERROR("Not compiled with GPU support"); AT_ERROR("Not compiled with GPU support");
#endif #endif
} else { } else {
CHECK_CPU(points);
CHECK_CPU(idxs);
CHECK_CPU(grad_zbuf);
CHECK_CPU(grad_dists);
return RasterizePointsBackwardCpu(points, idxs, grad_zbuf, grad_dists); return RasterizePointsBackwardCpu(points, idxs, grad_zbuf, grad_dists);
} }
} }

View File

@ -68,5 +68,9 @@ at::Tensor FarthestPointSampling(
AT_ERROR("Not compiled with GPU support."); AT_ERROR("Not compiled with GPU support.");
#endif #endif
} }
CHECK_CPU(points);
CHECK_CPU(lengths);
CHECK_CPU(K);
CHECK_CPU(start_idxs);
return FarthestPointSamplingCpu(points, lengths, K, start_idxs); return FarthestPointSamplingCpu(points, lengths, K, start_idxs);
} }

View File

@ -71,6 +71,8 @@ inline void SamplePdf(
AT_ERROR("Not compiled with GPU support."); AT_ERROR("Not compiled with GPU support.");
#endif #endif
} }
CHECK_CPU(weights);
CHECK_CPU(outputs);
CHECK_CONTIGUOUS(outputs); CHECK_CONTIGUOUS(outputs);
SamplePdfCpu(bins, weights, outputs, eps); SamplePdfCpu(bins, weights, outputs, eps);
} }