ignore cuda for cpu only installation

Summary:
Added if `WITH_CUDA` checks for points/mesh rasterization. If installing on cpu only then this causes `Undefined symbol` errors when trying to import pytorch3d.

We had these checks for all the other cuda files but not the rasterization files.

Thanks ppwwyyxx for the tip!

Reviewed By: ppwwyyxx, gkioxari

Differential Revision: D19801495

fbshipit-source-id: 20e7adccfdb33ac731c00a89414b2beaf0a35529
This commit is contained in:
Nikhila Ravi 2020-02-08 09:13:08 -08:00 committed by Facebook Github Bot
parent ca588a59d7
commit dcb094800f
2 changed files with 48 additions and 7 deletions

View File

@ -19,6 +19,7 @@ RasterizeMeshesNaiveCpu(
int faces_per_pixel, int faces_per_pixel,
bool perspective_correct); bool perspective_correct);
#ifdef WITH_CUDA
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor>
RasterizeMeshesNaiveCuda( RasterizeMeshesNaiveCuda(
const at::Tensor& face_verts, const at::Tensor& face_verts,
@ -28,7 +29,7 @@ RasterizeMeshesNaiveCuda(
float blur_radius, float blur_radius,
int num_closest, int num_closest,
bool perspective_correct); bool perspective_correct);
#endif
// Forward pass for rasterizing a batch of meshes. // Forward pass for rasterizing a batch of meshes.
// //
// Args: // Args:
@ -82,6 +83,7 @@ RasterizeMeshesNaive(
bool perspective_correct) { bool perspective_correct) {
// TODO: Better type checking. // TODO: Better type checking.
if (face_verts.type().is_cuda()) { if (face_verts.type().is_cuda()) {
#ifdef WITH_CUDA
return RasterizeMeshesNaiveCuda( return RasterizeMeshesNaiveCuda(
face_verts, face_verts,
mesh_to_face_first_idx, mesh_to_face_first_idx,
@ -90,6 +92,9 @@ RasterizeMeshesNaive(
blur_radius, blur_radius,
faces_per_pixel, faces_per_pixel,
perspective_correct); perspective_correct);
#else
AT_ERROR("Not compiled with GPU support");
#endif
} else { } else {
return RasterizeMeshesNaiveCpu( return RasterizeMeshesNaiveCpu(
face_verts, face_verts,
@ -114,6 +119,7 @@ torch::Tensor RasterizeMeshesBackwardCpu(
const torch::Tensor& grad_dists, const torch::Tensor& grad_dists,
bool perspective_correct); bool perspective_correct);
#ifdef WITH_CUDA
torch::Tensor RasterizeMeshesBackwardCuda( torch::Tensor RasterizeMeshesBackwardCuda(
const torch::Tensor& face_verts, const torch::Tensor& face_verts,
const torch::Tensor& pix_to_face, const torch::Tensor& pix_to_face,
@ -121,6 +127,7 @@ torch::Tensor RasterizeMeshesBackwardCuda(
const torch::Tensor& grad_zbuf, const torch::Tensor& grad_zbuf,
const torch::Tensor& grad_dists, const torch::Tensor& grad_dists,
bool perspective_correct); bool perspective_correct);
#endif
// Args: // Args:
// face_verts: float32 Tensor of shape (F, 3, 3) (from forward pass) giving // face_verts: float32 Tensor of shape (F, 3, 3) (from forward pass) giving
@ -154,6 +161,7 @@ torch::Tensor RasterizeMeshesBackward(
const torch::Tensor& grad_dists, const torch::Tensor& grad_dists,
bool perspective_correct) { bool perspective_correct) {
if (face_verts.type().is_cuda()) { if (face_verts.type().is_cuda()) {
#ifdef WITH_CUDA
return RasterizeMeshesBackwardCuda( return RasterizeMeshesBackwardCuda(
face_verts, face_verts,
pix_to_face, pix_to_face,
@ -161,6 +169,9 @@ torch::Tensor RasterizeMeshesBackward(
grad_bary, grad_bary,
grad_dists, grad_dists,
perspective_correct); perspective_correct);
#else
AT_ERROR("Not compiled with GPU support");
#endif
} else { } else {
return RasterizeMeshesBackwardCpu( return RasterizeMeshesBackwardCpu(
face_verts, face_verts,
@ -185,6 +196,7 @@ torch::Tensor RasterizeMeshesCoarseCpu(
int bin_size, int bin_size,
int max_faces_per_bin); int max_faces_per_bin);
#ifdef WITH_CUDA
torch::Tensor RasterizeMeshesCoarseCuda( torch::Tensor RasterizeMeshesCoarseCuda(
const torch::Tensor& face_verts, const torch::Tensor& face_verts,
const torch::Tensor& mesh_to_face_first_idx, const torch::Tensor& mesh_to_face_first_idx,
@ -193,7 +205,7 @@ torch::Tensor RasterizeMeshesCoarseCuda(
float blur_radius, float blur_radius,
int bin_size, int bin_size,
int max_faces_per_bin); int max_faces_per_bin);
#endif
// Args: // Args:
// face_verts: Tensor of shape (F, 3, 3) giving (packed) vertex positions for // face_verts: Tensor of shape (F, 3, 3) giving (packed) vertex positions for
// faces in all the meshes in the batch. Concretely, // faces in all the meshes in the batch. Concretely,
@ -225,6 +237,7 @@ torch::Tensor RasterizeMeshesCoarse(
int bin_size, int bin_size,
int max_faces_per_bin) { int max_faces_per_bin) {
if (face_verts.type().is_cuda()) { if (face_verts.type().is_cuda()) {
#ifdef WITH_CUDA
return RasterizeMeshesCoarseCuda( return RasterizeMeshesCoarseCuda(
face_verts, face_verts,
mesh_to_face_first_idx, mesh_to_face_first_idx,
@ -233,6 +246,9 @@ torch::Tensor RasterizeMeshesCoarse(
blur_radius, blur_radius,
bin_size, bin_size,
max_faces_per_bin); max_faces_per_bin);
#else
AT_ERROR("Not compiled with GPU support");
#endif
} else { } else {
return RasterizeMeshesCoarseCpu( return RasterizeMeshesCoarseCpu(
face_verts, face_verts,
@ -249,6 +265,7 @@ torch::Tensor RasterizeMeshesCoarse(
// * FINE RASTERIZATION * // * FINE RASTERIZATION *
// **************************************************************************** // ****************************************************************************
#ifdef WITH_CUDA
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
RasterizeMeshesFineCuda( RasterizeMeshesFineCuda(
const torch::Tensor& face_verts, const torch::Tensor& face_verts,
@ -258,7 +275,7 @@ RasterizeMeshesFineCuda(
int bin_size, int bin_size,
int faces_per_pixel, int faces_per_pixel,
bool perspective_correct); bool perspective_correct);
#endif
// Args: // Args:
// face_verts: Tensor of shape (F, 3, 3) giving (packed) vertex positions for // face_verts: Tensor of shape (F, 3, 3) giving (packed) vertex positions for
// faces in all the meshes in the batch. Concretely, // faces in all the meshes in the batch. Concretely,
@ -306,6 +323,7 @@ RasterizeMeshesFine(
int faces_per_pixel, int faces_per_pixel,
bool perspective_correct) { bool perspective_correct) {
if (face_verts.type().is_cuda()) { if (face_verts.type().is_cuda()) {
#ifdef WITH_CUDA
return RasterizeMeshesFineCuda( return RasterizeMeshesFineCuda(
face_verts, face_verts,
bin_faces, bin_faces,
@ -314,6 +332,9 @@ RasterizeMeshesFine(
bin_size, bin_size,
faces_per_pixel, faces_per_pixel,
perspective_correct); perspective_correct);
#else
AT_ERROR("Not compiled with GPU support");
#endif
} else { } else {
AT_ERROR("NOT IMPLEMENTED"); AT_ERROR("NOT IMPLEMENTED");
} }

View File

@ -15,13 +15,14 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> RasterizePointsNaiveCpu(
const float radius, const float radius,
const int points_per_pixel); const int points_per_pixel);
#ifdef WITH_CUDA
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
RasterizePointsNaiveCuda( RasterizePointsNaiveCuda(
const torch::Tensor& points, const torch::Tensor& points,
const int image_size, const int image_size,
const float radius, const float radius,
const int points_per_pixel); const int points_per_pixel);
#endif
// Naive (forward) pointcloud rasterization: For each pixel, for each point, // Naive (forward) pointcloud rasterization: For each pixel, for each point,
// check whether that point hits the pixel. // check whether that point hits the pixel.
// //
@ -47,8 +48,12 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> RasterizePointsNaive(
const float radius, const float radius,
const int points_per_pixel) { const int points_per_pixel) {
if (points.type().is_cuda()) { if (points.type().is_cuda()) {
#ifdef WITH_CUDA
return RasterizePointsNaiveCuda( return RasterizePointsNaiveCuda(
points, image_size, radius, points_per_pixel); points, image_size, radius, points_per_pixel);
#else
AT_ERROR("Not compiled with GPU support");
#endif
} else { } else {
return RasterizePointsNaiveCpu( return RasterizePointsNaiveCpu(
points, image_size, radius, points_per_pixel); points, image_size, radius, points_per_pixel);
@ -66,13 +71,14 @@ torch::Tensor RasterizePointsCoarseCpu(
const int bin_size, const int bin_size,
const int max_points_per_bin); const int max_points_per_bin);
#ifdef WITH_CUDA
torch::Tensor RasterizePointsCoarseCuda( torch::Tensor RasterizePointsCoarseCuda(
const torch::Tensor& points, const torch::Tensor& points,
const int image_size, const int image_size,
const float radius, const float radius,
const int bin_size, const int bin_size,
const int max_points_per_bin); const int max_points_per_bin);
#endif
// Args: // Args:
// points: Tensor of shape (N, P, 3) // points: Tensor of shape (N, P, 3)
// radius: Radius of points to rasterize (in NDC units) // radius: Radius of points to rasterize (in NDC units)
@ -91,8 +97,12 @@ torch::Tensor RasterizePointsCoarse(
const int bin_size, const int bin_size,
const int max_points_per_bin) { const int max_points_per_bin) {
if (points.type().is_cuda()) { if (points.type().is_cuda()) {
#ifdef WITH_CUDA
return RasterizePointsCoarseCuda( return RasterizePointsCoarseCuda(
points, image_size, radius, bin_size, max_points_per_bin); points, image_size, radius, bin_size, max_points_per_bin);
#else
AT_ERROR("Not compiled with GPU support");
#endif
} else { } else {
return RasterizePointsCoarseCpu( return RasterizePointsCoarseCpu(
points, image_size, radius, bin_size, max_points_per_bin); points, image_size, radius, bin_size, max_points_per_bin);
@ -103,6 +113,7 @@ torch::Tensor RasterizePointsCoarse(
// * FINE RASTERIZATION * // * FINE RASTERIZATION *
// **************************************************************************** // ****************************************************************************
#ifdef WITH_CUDA
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> RasterizePointsFineCuda( std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> RasterizePointsFineCuda(
const torch::Tensor& points, const torch::Tensor& points,
const torch::Tensor& bin_points, const torch::Tensor& bin_points,
@ -110,7 +121,7 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> RasterizePointsFineCuda(
const float radius, const float radius,
const int bin_size, const int bin_size,
const int points_per_pixel); const int points_per_pixel);
#endif
// Args: // Args:
// points: float32 Tensor of shape (N, P, 3) // points: float32 Tensor of shape (N, P, 3)
// bin_points: int32 Tensor of shape (N, B, B, M) giving the indices of points // bin_points: int32 Tensor of shape (N, B, B, M) giving the indices of points
@ -137,8 +148,12 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> RasterizePointsFine(
const int bin_size, const int bin_size,
const int points_per_pixel) { const int points_per_pixel) {
if (points.type().is_cuda()) { if (points.type().is_cuda()) {
#ifdef WITH_CUDA
return RasterizePointsFineCuda( return RasterizePointsFineCuda(
points, bin_points, image_size, radius, bin_size, points_per_pixel); points, bin_points, image_size, radius, bin_size, points_per_pixel);
#else
AT_ERROR("Not compiled with GPU support");
#endif
} else { } else {
AT_ERROR("NOT IMPLEMENTED"); AT_ERROR("NOT IMPLEMENTED");
} }
@ -154,12 +169,13 @@ torch::Tensor RasterizePointsBackwardCpu(
const torch::Tensor& grad_zbuf, const torch::Tensor& grad_zbuf,
const torch::Tensor& grad_dists); const torch::Tensor& grad_dists);
#ifdef WITH_CUDA
torch::Tensor RasterizePointsBackwardCuda( torch::Tensor RasterizePointsBackwardCuda(
const torch::Tensor& points, const torch::Tensor& points,
const torch::Tensor& idxs, const torch::Tensor& idxs,
const torch::Tensor& grad_zbuf, const torch::Tensor& grad_zbuf,
const torch::Tensor& grad_dists); const torch::Tensor& grad_dists);
#endif
// Args: // Args:
// points: float32 Tensor of shape (N, P, 3) // points: float32 Tensor of shape (N, P, 3)
// idxs: int32 Tensor of shape (N, H, W, K) (from forward pass) // idxs: int32 Tensor of shape (N, H, W, K) (from forward pass)
@ -178,7 +194,11 @@ torch::Tensor RasterizePointsBackward(
const torch::Tensor& grad_zbuf, const torch::Tensor& grad_zbuf,
const torch::Tensor& grad_dists) { const torch::Tensor& grad_dists) {
if (points.type().is_cuda()) { if (points.type().is_cuda()) {
#ifdef WITH_CUDA
return RasterizePointsBackwardCuda(points, idxs, grad_zbuf, grad_dists); return RasterizePointsBackwardCuda(points, idxs, grad_zbuf, grad_dists);
#else
AT_ERROR("Not compiled with GPU support");
#endif
} else { } else {
return RasterizePointsBackwardCpu(points, idxs, grad_zbuf, grad_dists); return RasterizePointsBackwardCpu(points, idxs, grad_zbuf, grad_dists);
} }