337 lines
10 KiB
C++
337 lines
10 KiB
C++
/**
|
|
* This file includes code originally from the lietorch repository:
|
|
* https://github.com/princeton-vl/lietorch
|
|
* Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
|
*/
|
|
|
|
#include <torch/extension.h>
|
|
#include <vector>
|
|
#include "lietorch_gpu.h"
|
|
#include "lietorch_cpu.h"
|
|
#include <c10/cuda/CUDAGuard.h>
|
|
|
|
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
|
|
|
|
/* Interface for cuda and c++ group operations
|
|
|
|
enum group_t { SO3=1, SE3=2, Sim3=3 };
|
|
X, Y, Z: (uppercase) Lie Group Elements
|
|
a, b, c: (lowercase) Lie Algebra Elements
|
|
*/
|
|
|
|
// Unary operations
|
|
torch::Tensor expm(int group_index, torch::Tensor a) {
|
|
CHECK_CONTIGUOUS(a);
|
|
if (a.device().type() == torch::DeviceType::CPU) {
|
|
return exp_forward_cpu(group_index, a);
|
|
|
|
} else if (a.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
|
|
return exp_forward_gpu(group_index, a);
|
|
}
|
|
|
|
return a;
|
|
}
|
|
|
|
std::vector<torch::Tensor> expm_backward(int group_index, torch::Tensor grad, torch::Tensor a) {
|
|
CHECK_CONTIGUOUS(a);
|
|
CHECK_CONTIGUOUS(grad);
|
|
if (a.device().type() == torch::DeviceType::CPU) {
|
|
return exp_backward_cpu(group_index, grad, a);
|
|
|
|
} else if (a.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
|
|
return exp_backward_gpu(group_index, grad, a);
|
|
}
|
|
|
|
return {};
|
|
}
|
|
|
|
torch::Tensor logm(int group_index, torch::Tensor X) {
|
|
CHECK_CONTIGUOUS(X);
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return log_forward_cpu(group_index, X);
|
|
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return log_forward_gpu(group_index, X);
|
|
}
|
|
|
|
return X;
|
|
}
|
|
|
|
std::vector<torch::Tensor> logm_backward(int group_index, torch::Tensor grad, torch::Tensor X) {
|
|
CHECK_CONTIGUOUS(X);
|
|
CHECK_CONTIGUOUS(grad);
|
|
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return log_backward_cpu(group_index, grad, X);
|
|
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return log_backward_gpu(group_index, grad, X);
|
|
}
|
|
|
|
return {};
|
|
}
|
|
|
|
torch::Tensor inv(int group_index, torch::Tensor X) {
|
|
CHECK_CONTIGUOUS(X);
|
|
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return inv_forward_cpu(group_index, X);
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return inv_forward_gpu(group_index, X);
|
|
}
|
|
|
|
return X;
|
|
}
|
|
|
|
std::vector<torch::Tensor> inv_backward(int group_index, torch::Tensor grad, torch::Tensor X) {
|
|
CHECK_CONTIGUOUS(X);
|
|
CHECK_CONTIGUOUS(grad);
|
|
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return inv_backward_cpu(group_index, grad, X);
|
|
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return inv_backward_gpu(group_index, grad, X);
|
|
}
|
|
|
|
return {};
|
|
}
|
|
|
|
// Binary operations
|
|
|
|
torch::Tensor mul(int group_index, torch::Tensor X, torch::Tensor Y) {
|
|
CHECK_CONTIGUOUS(X);
|
|
CHECK_CONTIGUOUS(Y);
|
|
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return mul_forward_cpu(group_index, X, Y);
|
|
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return mul_forward_gpu(group_index, X, Y);
|
|
}
|
|
|
|
return X;
|
|
}
|
|
|
|
std::vector<torch::Tensor> mul_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor Y) {
|
|
CHECK_CONTIGUOUS(X);
|
|
CHECK_CONTIGUOUS(Y);
|
|
CHECK_CONTIGUOUS(grad);
|
|
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return mul_backward_cpu(group_index, grad, X, Y);
|
|
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return mul_backward_gpu(group_index, grad, X, Y);
|
|
}
|
|
|
|
return {};
|
|
}
|
|
|
|
torch::Tensor adj(int group_index, torch::Tensor X, torch::Tensor a) {
|
|
CHECK_CONTIGUOUS(X);
|
|
CHECK_CONTIGUOUS(a);
|
|
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return adj_forward_cpu(group_index, X, a);
|
|
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return adj_forward_gpu(group_index, X, a);
|
|
}
|
|
|
|
return X;
|
|
}
|
|
|
|
std::vector<torch::Tensor> adj_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
|
|
CHECK_CONTIGUOUS(X);
|
|
CHECK_CONTIGUOUS(a);
|
|
CHECK_CONTIGUOUS(grad);
|
|
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return adj_backward_cpu(group_index, grad, X, a);
|
|
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return adj_backward_gpu(group_index, grad, X, a);
|
|
}
|
|
|
|
return {};
|
|
}
|
|
|
|
torch::Tensor adjT(int group_index, torch::Tensor X, torch::Tensor a) {
|
|
CHECK_CONTIGUOUS(X);
|
|
CHECK_CONTIGUOUS(a);
|
|
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return adjT_forward_cpu(group_index, X, a);
|
|
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return adjT_forward_gpu(group_index, X, a);
|
|
}
|
|
|
|
return X;
|
|
}
|
|
|
|
std::vector<torch::Tensor> adjT_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
|
|
CHECK_CONTIGUOUS(X);
|
|
CHECK_CONTIGUOUS(a);
|
|
CHECK_CONTIGUOUS(grad);
|
|
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return adjT_backward_cpu(group_index, grad, X, a);
|
|
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return adjT_backward_gpu(group_index, grad, X, a);
|
|
}
|
|
|
|
return {};
|
|
}
|
|
|
|
torch::Tensor act(int group_index, torch::Tensor X, torch::Tensor p) {
|
|
CHECK_CONTIGUOUS(X);
|
|
CHECK_CONTIGUOUS(p);
|
|
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return act_forward_cpu(group_index, X, p);
|
|
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return act_forward_gpu(group_index, X, p);
|
|
}
|
|
|
|
return X;
|
|
}
|
|
|
|
std::vector<torch::Tensor> act_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
|
|
CHECK_CONTIGUOUS(X);
|
|
CHECK_CONTIGUOUS(p);
|
|
CHECK_CONTIGUOUS(grad);
|
|
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return act_backward_cpu(group_index, grad, X, p);
|
|
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return act_backward_gpu(group_index, grad, X, p);
|
|
}
|
|
|
|
return {};
|
|
}
|
|
|
|
torch::Tensor act4(int group_index, torch::Tensor X, torch::Tensor p) {
|
|
CHECK_CONTIGUOUS(X);
|
|
CHECK_CONTIGUOUS(p);
|
|
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return act4_forward_cpu(group_index, X, p);
|
|
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return act4_forward_gpu(group_index, X, p);
|
|
}
|
|
|
|
return X;
|
|
}
|
|
|
|
std::vector<torch::Tensor> act4_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
|
|
CHECK_CONTIGUOUS(X);
|
|
CHECK_CONTIGUOUS(p);
|
|
CHECK_CONTIGUOUS(grad);
|
|
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return act4_backward_cpu(group_index, grad, X, p);
|
|
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return act4_backward_gpu(group_index, grad, X, p);
|
|
}
|
|
|
|
return {};
|
|
}
|
|
|
|
torch::Tensor projector(int group_index, torch::Tensor X) {
|
|
CHECK_CONTIGUOUS(X);
|
|
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return orthogonal_projector_cpu(group_index, X);
|
|
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return orthogonal_projector_gpu(group_index, X);
|
|
}
|
|
|
|
return X;
|
|
}
|
|
|
|
torch::Tensor as_matrix(int group_index, torch::Tensor X) {
|
|
CHECK_CONTIGUOUS(X);
|
|
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return as_matrix_forward_cpu(group_index, X);
|
|
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return as_matrix_forward_gpu(group_index, X);
|
|
}
|
|
|
|
return X;
|
|
}
|
|
|
|
torch::Tensor Jinv(int group_index, torch::Tensor X, torch::Tensor a) {
|
|
CHECK_CONTIGUOUS(X);
|
|
CHECK_CONTIGUOUS(a);
|
|
|
|
if (X.device().type() == torch::DeviceType::CPU) {
|
|
return jleft_forward_cpu(group_index, X, a);
|
|
|
|
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
|
return jleft_forward_gpu(group_index, X, a);
|
|
}
|
|
|
|
return a;
|
|
}
|
|
|
|
void pybind_lietorch_ext(py::module &m) {
|
|
// {exp, log, inv, mul, adj, adjT, act, act4} forward/backward bindings
|
|
m.def("expm", &expm, "exp map forward");
|
|
m.def("expm_backward", &expm_backward, "exp map backward");
|
|
|
|
m.def("logm", &logm, "log map forward");
|
|
m.def("logm_backward", &logm_backward, "log map backward");
|
|
|
|
m.def("inv", &inv, "inverse operator");
|
|
m.def("inv_backward", &inv_backward, "inverse operator backward");
|
|
|
|
m.def("mul", &mul, "group operator");
|
|
m.def("mul_backward", &mul_backward, "group operator backward");
|
|
|
|
m.def("adj", &adj, "adjoint operator");
|
|
m.def("adj_backward", &adj_backward, "adjoint operator backward");
|
|
|
|
m.def("adjT", &adjT, "transposed adjoint operator");
|
|
m.def("adjT_backward", &adjT_backward, "transposed adjoint operator backward");
|
|
|
|
m.def("act", &act, "action on point");
|
|
m.def("act_backward", &act_backward, "action on point backward");
|
|
|
|
m.def("act4", &act4, "action on homogeneous point");
|
|
m.def("act4_backward", &act4_backward, "action on homogeneous point backward");
|
|
|
|
// functions with no gradient
|
|
m.def("as_matrix", &as_matrix, "convert to matrix");
|
|
m.def("projector", &projector, "orthogonal projection matrix");
|
|
m.def("Jinv", &Jinv, "left inverse jacobian operator");
|
|
}
|