Skip to content

Commit

Permalink
CPU implementation for point_mesh functions
Browse files Browse the repository at this point in the history
Summary:
point_mesh functions were missing CPU implementations.
The indices returned are not always matching, possibly due to numerical instability.

Reviewed By: gkioxari

Differential Revision: D21594264

fbshipit-source-id: 3016930e2a9a0f3cd8b3ac4c94a92c9411c0989d
  • Loading branch information
bottler authored and facebook-github-bot committed Jun 15, 2020
1 parent 7f1e63a commit 74659ae
Show file tree
Hide file tree
Showing 6 changed files with 878 additions and 31 deletions.
398 changes: 398 additions & 0 deletions pytorch3d/csrc/point_mesh/point_mesh.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,398 @@
// Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

#include <torch/extension.h>
#include <array>
#include <limits>
#include "utils/geometry_utils.h"
#include "utils/vec3.h"

// - We start with implementations of simple operations on points, edges and
// faces. The hull of H points is a point if H=1, an edge if H=2, a face if H=3.

template <typename T>
vec3<T> ExtractPoint(const at::TensorAccessor<T, 1>& t) {
return vec3(t[0], t[1], t[2]);
}

template <class Accessor>
struct ExtractHullHelper {
template <int H>
static std::array<vec3<std::remove_pointer_t<typename Accessor::PtrType>>, H>
get(const Accessor& t);

template <>
static std::array<vec3<std::remove_pointer_t<typename Accessor::PtrType>>, 1>
get<1>(const Accessor& t) {
return {ExtractPoint(t)};
}

template <>
static std::array<vec3<std::remove_pointer_t<typename Accessor::PtrType>>, 2>
get<2>(const Accessor& t) {
return {ExtractPoint(t[0]), ExtractPoint(t[1])};
}

template <>
static std::array<vec3<std::remove_pointer_t<typename Accessor::PtrType>>, 3>
get<3>(const Accessor& t) {
return {ExtractPoint(t[0]), ExtractPoint(t[1]), ExtractPoint(t[2])};
}
};

template <int H, typename Accessor>
std::array<vec3<std::remove_pointer_t<typename Accessor::PtrType>>, H>
ExtractHull(const Accessor& t) {
return ExtractHullHelper<Accessor>::template get<H>(t);
}

template <typename T>
void IncrementPoint(at::TensorAccessor<T, 1>&& t, const vec3<T>& point) {
t[0] += point.x;
t[1] += point.y;
t[2] += point.z;
}

// distance between the convex hull of A points and B points
// this could be done in c++17 with tuple_cat and invoke
template <typename T>
T HullDistance(
const std::array<vec3<T>, 1>& a,
const std::array<vec3<T>, 2>& b) {
using std::get;
return PointLine3DistanceForward(get<0>(a), get<0>(b), get<1>(b));
}
template <typename T>
T HullDistance(
const std::array<vec3<T>, 1>& a,
const std::array<vec3<T>, 3>& b) {
using std::get;
return PointTriangle3DistanceForward(
get<0>(a), get<0>(b), get<1>(b), get<2>(b));
}
template <typename T>
T HullDistance(
const std::array<vec3<T>, 2>& a,
const std::array<vec3<T>, 1>& b) {
return HullDistance(b, a);
}
template <typename T>
T HullDistance(
const std::array<vec3<T>, 3>& a,
const std::array<vec3<T>, 1>& b) {
return HullDistance(b, a);
}

template <typename T>
void HullHullDistanceBackward(
const std::array<vec3<T>, 1>& a,
const std::array<vec3<T>, 2>& b,
T grad_dist,
at::TensorAccessor<T, 1>&& grad_a,
at::TensorAccessor<T, 2>&& grad_b) {
using std::get;
auto res =
PointLine3DistanceBackward(get<0>(a), get<0>(b), get<1>(b), grad_dist);
IncrementPoint(std::move(grad_a), get<0>(res));
IncrementPoint(grad_b[0], get<1>(res));
IncrementPoint(grad_b[1], get<2>(res));
}
template <typename T>
void HullHullDistanceBackward(
const std::array<vec3<T>, 1>& a,
const std::array<vec3<T>, 3>& b,
T grad_dist,
at::TensorAccessor<T, 1>&& grad_a,
at::TensorAccessor<T, 2>&& grad_b) {
using std::get;
auto res = PointTriangle3DistanceBackward(
get<0>(a), get<0>(b), get<1>(b), get<2>(b), grad_dist);
IncrementPoint(std::move(grad_a), get<0>(res));
IncrementPoint(grad_b[0], get<1>(res));
IncrementPoint(grad_b[1], get<2>(res));
IncrementPoint(grad_b[2], get<3>(res));
}
template <typename T>
void HullHullDistanceBackward(
const std::array<vec3<T>, 3>& a,
const std::array<vec3<T>, 1>& b,
T grad_dist,
at::TensorAccessor<T, 2>&& grad_a,
at::TensorAccessor<T, 1>&& grad_b) {
return HullHullDistanceBackward(
b, a, grad_dist, std::move(grad_b), std::move(grad_a));
}
template <typename T>
void HullHullDistanceBackward(
const std::array<vec3<T>, 2>& a,
const std::array<vec3<T>, 1>& b,
T grad_dist,
at::TensorAccessor<T, 2>&& grad_a,
at::TensorAccessor<T, 1>&& grad_b) {
return HullHullDistanceBackward(
b, a, grad_dist, std::move(grad_b), std::move(grad_a));
}

template <int H>
void ValidateShape(const at::Tensor& as) {
if (H == 1) {
TORCH_CHECK(as.size(1) == 3);
} else {
TORCH_CHECK(as.size(2) == 3 && as.size(1) == H);
}
}

// ----------- Here begins the implementation of each top-level
// function using non-type template parameters to
// implement all the cases in one go. ----------- //

template <int H1, int H2>
std::tuple<at::Tensor, at::Tensor> HullHullDistanceForwardCpu(
const at::Tensor& as,
const at::Tensor& as_first_idx,
const at::Tensor& bs,
const at::Tensor& bs_first_idx) {
const int64_t A_N = as.size(0);
const int64_t B_N = bs.size(0);
const int64_t BATCHES = as_first_idx.size(0);

ValidateShape<H1>(as);
ValidateShape<H2>(bs);

TORCH_CHECK(bs_first_idx.size(0) == BATCHES);

// clang-format off
at::Tensor dists = at::zeros({A_N,}, as.options());
at::Tensor idxs = at::zeros({A_N,}, as_first_idx.options());
// clang-format on

auto as_a = as.accessor < float, H1 == 1 ? 2 : 3 > ();
auto bs_a = bs.accessor < float, H2 == 1 ? 2 : 3 > ();
auto as_first_idx_a = as_first_idx.accessor<int64_t, 1>();
auto bs_first_idx_a = bs_first_idx.accessor<int64_t, 1>();
auto dists_a = dists.accessor<float, 1>();
auto idxs_a = idxs.accessor<int64_t, 1>();
int64_t a_batch_end = 0;
int64_t b_batch_start = 0, b_batch_end = 0;
int64_t batch_idx = 0;
for (int64_t a_n = 0; a_n < A_N; ++a_n) {
if (a_n == a_batch_end) {
++batch_idx;
b_batch_start = b_batch_end;
if (batch_idx == BATCHES) {
a_batch_end = std::numeric_limits<int64_t>::max();
b_batch_end = B_N;
} else {
a_batch_end = as_first_idx_a[batch_idx];
b_batch_end = bs_first_idx_a[batch_idx];
}
}
float min_dist = std::numeric_limits<float>::max();
size_t min_idx = 0;
auto a = ExtractHull<H1>(as_a[a_n]);
for (int64_t b_n = b_batch_start; b_n < b_batch_end; ++b_n) {
float dist = HullDistance(a, ExtractHull<H2>(bs_a[b_n]));
if (dist <= min_dist) {
min_dist = dist;
min_idx = b_n;
}
}
dists_a[a_n] = min_dist;
idxs_a[a_n] = min_idx;
}

return std::make_tuple(dists, idxs);
}

template <int H1, int H2>
std::tuple<at::Tensor, at::Tensor> HullHullDistanceBackwardCpu(
const at::Tensor& as,
const at::Tensor& bs,
const at::Tensor& idx_bs,
const at::Tensor& grad_dists) {
const int64_t A_N = as.size(0);

TORCH_CHECK(idx_bs.size(0) == A_N);
TORCH_CHECK(grad_dists.size(0) == A_N);
ValidateShape<H1>(as);
ValidateShape<H2>(bs);

at::Tensor grad_as = at::zeros_like(as);
at::Tensor grad_bs = at::zeros_like(bs);

auto as_a = as.accessor < float, H1 == 1 ? 2 : 3 > ();
auto bs_a = bs.accessor < float, H2 == 1 ? 2 : 3 > ();
auto grad_as_a = grad_as.accessor < float, H1 == 1 ? 2 : 3 > ();
auto grad_bs_a = grad_bs.accessor < float, H2 == 1 ? 2 : 3 > ();
auto idx_bs_a = idx_bs.accessor<int64_t, 1>();
auto grad_dists_a = grad_dists.accessor<float, 1>();

for (int64_t a_n = 0; a_n < A_N; ++a_n) {
auto a = ExtractHull<H1>(as_a[a_n]);
auto b = ExtractHull<H2>(bs_a[idx_bs_a[a_n]]);
HullHullDistanceBackward(
a, b, grad_dists_a[a_n], grad_as_a[a_n], grad_bs_a[idx_bs_a[a_n]]);
}
return std::make_tuple(grad_as, grad_bs);
}

template <int H>
torch::Tensor PointHullArrayDistanceForwardCpu(
const torch::Tensor& points,
const torch::Tensor& bs) {
const int64_t P = points.size(0);
const int64_t B_N = bs.size(0);

TORCH_CHECK(points.size(1) == 3, "points must be of shape Px3");
ValidateShape<H>(bs);

at::Tensor dists = at::zeros({P, B_N}, points.options());
auto points_a = points.accessor<float, 2>();
auto bs_a = bs.accessor<float, 3>();
auto dists_a = dists.accessor<float, 2>();
for (int64_t p = 0; p < P; ++p) {
auto point = ExtractHull<1>(points_a[p]);
auto dest = dists_a[p];
for (int64_t b_n = 0; b_n < B_N; ++b_n) {
auto b = ExtractHull<H>(bs_a[b_n]);
dest[b_n] = HullDistance(point, b);
}
}
return dists;
}

template <int H>
std::tuple<at::Tensor, at::Tensor> PointHullArrayDistanceBackwardCpu(
const at::Tensor& points,
const at::Tensor& bs,
const at::Tensor& grad_dists) {
const int64_t P = points.size(0);
const int64_t B_N = bs.size(0);

TORCH_CHECK(points.size(1) == 3, "points must be of shape Px3");
ValidateShape<H>(bs);
TORCH_CHECK((grad_dists.size(0) == P) && (grad_dists.size(1) == B_N));

at::Tensor grad_points = at::zeros({P, 3}, points.options());
at::Tensor grad_bs = at::zeros({B_N, H, 3}, bs.options());

auto points_a = points.accessor<float, 2>();
auto bs_a = bs.accessor<float, 3>();
auto grad_dists_a = grad_dists.accessor<float, 2>();
auto grad_points_a = grad_points.accessor<float, 2>();
auto grad_bs_a = grad_bs.accessor<float, 3>();
for (int64_t p = 0; p < P; ++p) {
auto point = ExtractHull<1>(points_a[p]);
auto grad_point = grad_points_a[p];
auto grad_dist = grad_dists_a[p];
for (int64_t b_n = 0; b_n < B_N; ++b_n) {
auto b = ExtractHull<H>(bs_a[b_n]);
HullHullDistanceBackward(
point, b, grad_dist[b_n], std::move(grad_point), grad_bs_a[b_n]);
}
}
return std::make_tuple(grad_points, grad_bs);
}

// ---------- Here begin the exported functions ------------ //

std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceForwardCpu(
const torch::Tensor& points,
const torch::Tensor& points_first_idx,
const torch::Tensor& tris,
const torch::Tensor& tris_first_idx) {
return HullHullDistanceForwardCpu<1, 3>(
points, points_first_idx, tris, tris_first_idx);
}

std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceBackwardCpu(
const torch::Tensor& points,
const torch::Tensor& tris,
const torch::Tensor& idx_points,
const torch::Tensor& grad_dists) {
return HullHullDistanceBackwardCpu<1, 3>(
points, tris, idx_points, grad_dists);
}

std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceForwardCpu(
const torch::Tensor& points,
const torch::Tensor& points_first_idx,
const torch::Tensor& tris,
const torch::Tensor& tris_first_idx) {
return HullHullDistanceForwardCpu<3, 1>(
tris, tris_first_idx, points, points_first_idx);
}

std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceBackwardCpu(
const torch::Tensor& points,
const torch::Tensor& tris,
const torch::Tensor& idx_tris,
const torch::Tensor& grad_dists) {
auto res =
HullHullDistanceBackwardCpu<3, 1>(tris, points, idx_tris, grad_dists);
return std::make_tuple(std::get<1>(res), std::get<0>(res));
}

torch::Tensor PointEdgeArrayDistanceForwardCpu(
const torch::Tensor& points,
const torch::Tensor& segms) {
return PointHullArrayDistanceForwardCpu<2>(points, segms);
}

std::tuple<at::Tensor, at::Tensor> PointFaceArrayDistanceBackwardCpu(
const at::Tensor& points,
const at::Tensor& tris,
const at::Tensor& grad_dists) {
return PointHullArrayDistanceBackwardCpu<3>(points, tris, grad_dists);
}

torch::Tensor PointFaceArrayDistanceForwardCpu(
const torch::Tensor& points,
const torch::Tensor& tris) {
return PointHullArrayDistanceForwardCpu<3>(points, tris);
}

std::tuple<at::Tensor, at::Tensor> PointEdgeArrayDistanceBackwardCpu(
const at::Tensor& points,
const at::Tensor& segms,
const at::Tensor& grad_dists) {
return PointHullArrayDistanceBackwardCpu<2>(points, segms, grad_dists);
}

std::tuple<torch::Tensor, torch::Tensor> PointEdgeDistanceForwardCpu(
const torch::Tensor& points,
const torch::Tensor& points_first_idx,
const torch::Tensor& segms,
const torch::Tensor& segms_first_idx,
const int64_t /*max_points*/) {
return HullHullDistanceForwardCpu<1, 2>(
points, points_first_idx, segms, segms_first_idx);
}

std::tuple<torch::Tensor, torch::Tensor> PointEdgeDistanceBackwardCpu(
const torch::Tensor& points,
const torch::Tensor& segms,
const torch::Tensor& idx_points,
const torch::Tensor& grad_dists) {
return HullHullDistanceBackwardCpu<1, 2>(
points, segms, idx_points, grad_dists);
}

std::tuple<torch::Tensor, torch::Tensor> EdgePointDistanceForwardCpu(
const torch::Tensor& points,
const torch::Tensor& points_first_idx,
const torch::Tensor& segms,
const torch::Tensor& segms_first_idx,
const int64_t /*max_segms*/) {
return HullHullDistanceForwardCpu<2, 1>(
segms, segms_first_idx, points, points_first_idx);
}

std::tuple<torch::Tensor, torch::Tensor> EdgePointDistanceBackwardCpu(
const torch::Tensor& points,
const torch::Tensor& segms,
const torch::Tensor& idx_segms,
const torch::Tensor& grad_dists) {
auto res =
HullHullDistanceBackwardCpu<2, 1>(segms, points, idx_segms, grad_dists);
return std::make_tuple(std::get<1>(res), std::get<0>(res));
}
Loading

0 comments on commit 74659ae

Please sign in to comment.