Add Auto Build Pipeline & New Wheels for windows
Package build based on: python 3.10 - 3.12 cuda 11.8 and 12.1
This commit is contained in:
Mr. For Example
+4
-1
@@ -164,4 +164,7 @@ nvdiffrast/
|
||||
*.safetensors
|
||||
|
||||
# Ignore test outputs
|
||||
outputs/
|
||||
outputs/
|
||||
|
||||
# Ignore git cloned project only for build dependencies
|
||||
_Pre_Builds/_Build_Dependencies
|
||||
+1
-2
@@ -62,8 +62,7 @@ RUN pip install --no-cache -r requirements.txt
|
||||
|
||||
WORKDIR /app/custom_nodes/ComfyUI-3D-Pack/
|
||||
COPY --chown=user:user requirements.txt requirements_post.txt ./
|
||||
COPY --chown=user:user Libs/simple-knn/ Libs/simple-knn/
|
||||
COPY --chown=user:user Gen_3D_Modules/TriplaneGaussian/ Gen_3D_Modules/TriplaneGaussian/
|
||||
COPY --chown=user:user _Pre_Builds/_Libs _Pre_Builds/_Libs
|
||||
RUN pip install --no-cache -r requirements.txt \
|
||||
# post requirements installation require gpu, setup
|
||||
# `nvidia-container-runtime`, for docker, see
|
||||
|
||||
@@ -0,0 +1,140 @@
|
||||
import sys
|
||||
import os
|
||||
from os.path import dirname
|
||||
import glob
|
||||
import platform
|
||||
import subprocess
|
||||
import re
|
||||
|
||||
SCRIPT_ROOT_ABS_PATH = dirname(os.path.abspath(__file__))
|
||||
DEPENDENCIES_FILE_ABS_PATH = os.path.join(SCRIPT_ROOT_ABS_PATH, "dependencies.txt")
|
||||
BUILD_REQUIREMENTS_FILE_ABS_PATH = os.path.join(SCRIPT_ROOT_ABS_PATH, "build_requirements.txt")
|
||||
BUILD_ROOT_ABS_PATH = dirname(SCRIPT_ROOT_ABS_PATH)
|
||||
DEPENDENCIES_ROOT_ABS_PATH = os.path.join(BUILD_ROOT_ABS_PATH, "_Build_Dependencies")
|
||||
LIBS_ROOT_ABS_PATH = os.path.join(BUILD_ROOT_ABS_PATH, "_Libs")
|
||||
|
||||
def get_parent_dirpath_n_level_up(abs_path, n=1):
|
||||
for i in range(n):
|
||||
abs_path = dirname(abs_path)
|
||||
return abs_path
|
||||
|
||||
COMFY_PYTHON_ABS_PATH = os.path.join(get_parent_dirpath_n_level_up(BUILD_ROOT_ABS_PATH, 4), "python_embeded", "python")
|
||||
|
||||
def get_platform_config_name(cuda_version):
|
||||
platform_config_name = "_Wheels"
|
||||
|
||||
# Add OS Type
|
||||
if platform.system() == 'Windows':
|
||||
platform_config_name += "_win"
|
||||
elif platform.system() == "Linux":
|
||||
platform_config_name += "_linux"
|
||||
else:
|
||||
raise NotImplementedError(f"Platform {platform.system()} not supported!")
|
||||
|
||||
# Add Python Version, only first two version numbers, e.g. 3.12.4 -> 312
|
||||
platform_config_name += "_py" + "".join(platform.python_version().split('.')[:-1])
|
||||
|
||||
# Add CUDA Version
|
||||
platform_config_name += "_" + cuda_version
|
||||
|
||||
return platform_config_name
|
||||
|
||||
def read_dependencies(file_path):
|
||||
dependencies = [ f.path for f in os.scandir(LIBS_ROOT_ABS_PATH) if f.is_dir() ]
|
||||
with open(file_path, "r") as f:
|
||||
dependencies += [line.strip() for line in f]
|
||||
|
||||
return dependencies
|
||||
|
||||
def get_dependency_dir(dependency):
|
||||
is_url = dependency[:6] == "https:"
|
||||
if is_url:
|
||||
dependency_name = re.split(r'[/\.]+', dependency)[-2]
|
||||
dependency_dir = os.path.join(DEPENDENCIES_ROOT_ABS_PATH, dependency_name)
|
||||
else:
|
||||
dependency_dir = os.path.join(LIBS_ROOT_ABS_PATH, dependency)
|
||||
return dependency_dir, is_url
|
||||
|
||||
def setup_build_env(python_path, args):
|
||||
# Set CMake argumens fo build packages based on CUDA
|
||||
os.environ["CMAKE_ARGS"]="-DBUILD_opencv_world=ON -DWITH_CUDA=ON -DCUDA_FAST_MATH=ON -DWITH_CUBLAS=ON -DCUDA_ARCH_PTX=9.0 -DWITH_NVCUVID=ON"
|
||||
subprocess.run([python_path, "-s", "-m", "pip", "install", "-r", BUILD_REQUIREMENTS_FILE_ABS_PATH])
|
||||
subprocess.run([
|
||||
python_path, "-s", "-m", "pip", "install",
|
||||
f"torch=={args.torch_version}", f"torchvision=={args.torchvision_version}",
|
||||
"--index-url", f"https://download.pytorch.org/whl/{args.cuda_version}"
|
||||
])
|
||||
|
||||
def clone_or_update_read_dependency(dependency_url, dependency_dir):
|
||||
if os.path.exists(dependency_dir):
|
||||
# Dependency already exists, update it
|
||||
subprocess.run(["git", "pull"], cwd=dependency_dir)
|
||||
else:
|
||||
# Clone the dependency
|
||||
subprocess.run(["git", "clone", "--recursive", dependency_url, dependency_dir])
|
||||
|
||||
def build_python_wheel(python_path, dependency_dir, output_dir):
|
||||
# Build wheel and move the wheel file we just built to the output directory
|
||||
subprocess.run([python_path, "setup.py", "bdist_wheel", "--dist-dir", output_dir], cwd=dependency_dir)
|
||||
|
||||
def main(args):
|
||||
# Get python executable from path
|
||||
python_path = sys.executable
|
||||
|
||||
# Create the output folder if it doesn't exist
|
||||
output_root_path = args.output_root_dir
|
||||
if output_root_path is None:
|
||||
output_root_path = os.path.join(BUILD_ROOT_ABS_PATH, get_platform_config_name(args.cuda_version))
|
||||
elif not os.path.isabs(output_root_path):
|
||||
output_root_path = os.path.join(BUILD_ROOT_ABS_PATH, output_root_path)
|
||||
|
||||
os.makedirs(output_root_path, exist_ok=True)
|
||||
|
||||
# Read dependencies names from the text file
|
||||
dependencies = read_dependencies(DEPENDENCIES_FILE_ABS_PATH)
|
||||
print(dependencies)
|
||||
|
||||
# Setup environment variables and fundamental packages for build
|
||||
setup_build_env(python_path, args)
|
||||
|
||||
# Build all dependencies and move them to output_root_path
|
||||
for dependency in dependencies:
|
||||
dependency_dir, is_url = get_dependency_dir(dependency)
|
||||
if is_url:
|
||||
clone_or_update_read_dependency(dependency, dependency_dir)
|
||||
|
||||
build_python_wheel(python_path, dependency_dir, output_root_path)
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--output_root_dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to the target build directory",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cuda_version",
|
||||
type=str,
|
||||
default="cu121",
|
||||
choices="[cu121, cu118]",
|
||||
help="CUDA Toolkit version",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--torch_version",
|
||||
type=str,
|
||||
default="2.3.0",
|
||||
help="Pytorch version",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--torchvision_version",
|
||||
type=str,
|
||||
default="0.18.0",
|
||||
help="Pytorch vision library version",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
main(args)
|
||||
|
||||
#C:\Users\reall\Softwares\Miniconda\envs\Comfy3D_Builds_py310_cu118\python _Pre_Builds\_Build_Scripts\auto_build_all.py --cuda_version cu118
|
||||
@@ -0,0 +1,2 @@
|
||||
cmake
|
||||
ninja
|
||||
@@ -0,0 +1,3 @@
|
||||
https://github.com/ashawkey/diff-gaussian-rasterization.git
|
||||
https://github.com/NVlabs/nvdiffrast.git
|
||||
https://github.com/ashawkey/kiuikit.git
|
||||
@@ -0,0 +1,3 @@
|
||||
import pointnet2_ops.pointnet2_modules
|
||||
import pointnet2_ops.pointnet2_utils
|
||||
from pointnet2_ops._version import __version__
|
||||
@@ -0,0 +1,5 @@
|
||||
#pragma once
|
||||
#include <torch/extension.h>
|
||||
|
||||
at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius,
|
||||
const int nsample);
|
||||
@@ -0,0 +1,41 @@
|
||||
#ifndef _CUDA_UTILS_H
|
||||
#define _CUDA_UTILS_H
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cmath>
|
||||
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <vector>
|
||||
|
||||
#define TOTAL_THREADS 512
|
||||
|
||||
inline int opt_n_threads(int work_size) {
|
||||
const int pow_2 = std::log(static_cast<double>(work_size)) / std::log(2.0);
|
||||
|
||||
return max(min(1 << pow_2, TOTAL_THREADS), 1);
|
||||
}
|
||||
|
||||
inline dim3 opt_block_config(int x, int y) {
|
||||
const int x_threads = opt_n_threads(x);
|
||||
const int y_threads =
|
||||
max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1);
|
||||
dim3 block_config(x_threads, y_threads, 1);
|
||||
|
||||
return block_config;
|
||||
}
|
||||
|
||||
#define CUDA_CHECK_ERRORS() \
|
||||
do { \
|
||||
cudaError_t err = cudaGetLastError(); \
|
||||
if (cudaSuccess != err) { \
|
||||
fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \
|
||||
cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \
|
||||
__FILE__); \
|
||||
exit(-1); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,5 @@
|
||||
#pragma once
|
||||
#include <torch/extension.h>
|
||||
|
||||
at::Tensor group_points(at::Tensor points, at::Tensor idx);
|
||||
at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n);
|
||||
@@ -0,0 +1,10 @@
|
||||
#pragma once
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <vector>
|
||||
|
||||
std::vector<at::Tensor> three_nn(at::Tensor unknowns, at::Tensor knows);
|
||||
at::Tensor three_interpolate(at::Tensor points, at::Tensor idx,
|
||||
at::Tensor weight);
|
||||
at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx,
|
||||
at::Tensor weight, const int m);
|
||||
@@ -0,0 +1,6 @@
|
||||
#pragma once
|
||||
#include <torch/extension.h>
|
||||
|
||||
at::Tensor gather_points(at::Tensor points, at::Tensor idx);
|
||||
at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, const int n);
|
||||
at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples);
|
||||
@@ -0,0 +1,25 @@
|
||||
#pragma once
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <torch/extension.h>
|
||||
|
||||
#define CHECK_CUDA(x) \
|
||||
do { \
|
||||
AT_ASSERT(x.is_cuda(), #x " must be a CUDA tensor"); \
|
||||
} while (0)
|
||||
|
||||
#define CHECK_CONTIGUOUS(x) \
|
||||
do { \
|
||||
AT_ASSERT(x.is_contiguous(), #x " must be a contiguous tensor"); \
|
||||
} while (0)
|
||||
|
||||
#define CHECK_IS_INT(x) \
|
||||
do { \
|
||||
AT_ASSERT(x.scalar_type() == at::ScalarType::Int, \
|
||||
#x " must be an int tensor"); \
|
||||
} while (0)
|
||||
|
||||
#define CHECK_IS_FLOAT(x) \
|
||||
do { \
|
||||
AT_ASSERT(x.scalar_type() == at::ScalarType::Float, \
|
||||
#x " must be a float tensor"); \
|
||||
} while (0)
|
||||
@@ -0,0 +1,32 @@
|
||||
#include "ball_query.h"
|
||||
#include "utils.h"
|
||||
|
||||
void query_ball_point_kernel_wrapper(int b, int n, int m, float radius,
|
||||
int nsample, const float *new_xyz,
|
||||
const float *xyz, int *idx);
|
||||
|
||||
at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius,
|
||||
const int nsample) {
|
||||
CHECK_CONTIGUOUS(new_xyz);
|
||||
CHECK_CONTIGUOUS(xyz);
|
||||
CHECK_IS_FLOAT(new_xyz);
|
||||
CHECK_IS_FLOAT(xyz);
|
||||
|
||||
if (new_xyz.is_cuda()) {
|
||||
CHECK_CUDA(xyz);
|
||||
}
|
||||
|
||||
at::Tensor idx =
|
||||
torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample},
|
||||
at::device(new_xyz.device()).dtype(at::ScalarType::Int));
|
||||
|
||||
if (new_xyz.is_cuda()) {
|
||||
query_ball_point_kernel_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1),
|
||||
radius, nsample, new_xyz.data_ptr<float>(),
|
||||
xyz.data_ptr<float>(), idx.data_ptr<int>());
|
||||
} else {
|
||||
AT_ASSERT(false, "CPU not supported");
|
||||
}
|
||||
|
||||
return idx;
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
#include <math.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
#include "cuda_utils.h"
|
||||
|
||||
// input: new_xyz(b, m, 3) xyz(b, n, 3)
|
||||
// output: idx(b, m, nsample)
|
||||
__global__ void query_ball_point_kernel(int b, int n, int m, float radius,
|
||||
int nsample,
|
||||
const float *__restrict__ new_xyz,
|
||||
const float *__restrict__ xyz,
|
||||
int *__restrict__ idx) {
|
||||
int batch_index = blockIdx.x;
|
||||
xyz += batch_index * n * 3;
|
||||
new_xyz += batch_index * m * 3;
|
||||
idx += m * nsample * batch_index;
|
||||
|
||||
int index = threadIdx.x;
|
||||
int stride = blockDim.x;
|
||||
|
||||
float radius2 = radius * radius;
|
||||
for (int j = index; j < m; j += stride) {
|
||||
float new_x = new_xyz[j * 3 + 0];
|
||||
float new_y = new_xyz[j * 3 + 1];
|
||||
float new_z = new_xyz[j * 3 + 2];
|
||||
for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) {
|
||||
float x = xyz[k * 3 + 0];
|
||||
float y = xyz[k * 3 + 1];
|
||||
float z = xyz[k * 3 + 2];
|
||||
float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) +
|
||||
(new_z - z) * (new_z - z);
|
||||
if (d2 < radius2) {
|
||||
if (cnt == 0) {
|
||||
for (int l = 0; l < nsample; ++l) {
|
||||
idx[j * nsample + l] = k;
|
||||
}
|
||||
}
|
||||
idx[j * nsample + cnt] = k;
|
||||
++cnt;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void query_ball_point_kernel_wrapper(int b, int n, int m, float radius,
|
||||
int nsample, const float *new_xyz,
|
||||
const float *xyz, int *idx) {
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
query_ball_point_kernel<<<b, opt_n_threads(m), 0, stream>>>(
|
||||
b, n, m, radius, nsample, new_xyz, xyz, idx);
|
||||
|
||||
CUDA_CHECK_ERRORS();
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
#include "ball_query.h"
|
||||
#include "group_points.h"
|
||||
#include "interpolate.h"
|
||||
#include "sampling.h"
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("gather_points", &gather_points);
|
||||
m.def("gather_points_grad", &gather_points_grad);
|
||||
m.def("furthest_point_sampling", &furthest_point_sampling);
|
||||
|
||||
m.def("three_nn", &three_nn);
|
||||
m.def("three_interpolate", &three_interpolate);
|
||||
m.def("three_interpolate_grad", &three_interpolate_grad);
|
||||
|
||||
m.def("ball_query", &ball_query);
|
||||
|
||||
m.def("group_points", &group_points);
|
||||
m.def("group_points_grad", &group_points_grad);
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
#include "group_points.h"
|
||||
#include "utils.h"
|
||||
|
||||
void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample,
|
||||
const float *points, const int *idx,
|
||||
float *out);
|
||||
|
||||
void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints,
|
||||
int nsample, const float *grad_out,
|
||||
const int *idx, float *grad_points);
|
||||
|
||||
at::Tensor group_points(at::Tensor points, at::Tensor idx) {
|
||||
CHECK_CONTIGUOUS(points);
|
||||
CHECK_CONTIGUOUS(idx);
|
||||
CHECK_IS_FLOAT(points);
|
||||
CHECK_IS_INT(idx);
|
||||
|
||||
if (points.is_cuda()) {
|
||||
CHECK_CUDA(idx);
|
||||
}
|
||||
|
||||
at::Tensor output =
|
||||
torch::zeros({points.size(0), points.size(1), idx.size(1), idx.size(2)},
|
||||
at::device(points.device()).dtype(at::ScalarType::Float));
|
||||
|
||||
if (points.is_cuda()) {
|
||||
group_points_kernel_wrapper(points.size(0), points.size(1), points.size(2),
|
||||
idx.size(1), idx.size(2),
|
||||
points.data_ptr<float>(), idx.data_ptr<int>(),
|
||||
output.data_ptr<float>());
|
||||
} else {
|
||||
AT_ASSERT(false, "CPU not supported");
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n) {
|
||||
CHECK_CONTIGUOUS(grad_out);
|
||||
CHECK_CONTIGUOUS(idx);
|
||||
CHECK_IS_FLOAT(grad_out);
|
||||
CHECK_IS_INT(idx);
|
||||
|
||||
if (grad_out.is_cuda()) {
|
||||
CHECK_CUDA(idx);
|
||||
}
|
||||
|
||||
at::Tensor output =
|
||||
torch::zeros({grad_out.size(0), grad_out.size(1), n},
|
||||
at::device(grad_out.device()).dtype(at::ScalarType::Float));
|
||||
|
||||
if (grad_out.is_cuda()) {
|
||||
group_points_grad_kernel_wrapper(
|
||||
grad_out.size(0), grad_out.size(1), n, idx.size(1), idx.size(2),
|
||||
grad_out.data_ptr<float>(), idx.data_ptr<int>(),
|
||||
output.data_ptr<float>());
|
||||
} else {
|
||||
AT_ASSERT(false, "CPU not supported");
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
#include "cuda_utils.h"
|
||||
|
||||
// input: points(b, c, n) idx(b, npoints, nsample)
|
||||
// output: out(b, c, npoints, nsample)
|
||||
__global__ void group_points_kernel(int b, int c, int n, int npoints,
|
||||
int nsample,
|
||||
const float *__restrict__ points,
|
||||
const int *__restrict__ idx,
|
||||
float *__restrict__ out) {
|
||||
int batch_index = blockIdx.x;
|
||||
points += batch_index * n * c;
|
||||
idx += batch_index * npoints * nsample;
|
||||
out += batch_index * npoints * nsample * c;
|
||||
|
||||
const int index = threadIdx.y * blockDim.x + threadIdx.x;
|
||||
const int stride = blockDim.y * blockDim.x;
|
||||
for (int i = index; i < c * npoints; i += stride) {
|
||||
const int l = i / npoints;
|
||||
const int j = i % npoints;
|
||||
for (int k = 0; k < nsample; ++k) {
|
||||
int ii = idx[j * nsample + k];
|
||||
out[(l * npoints + j) * nsample + k] = points[l * n + ii];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample,
|
||||
const float *points, const int *idx,
|
||||
float *out) {
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
group_points_kernel<<<b, opt_block_config(npoints, c), 0, stream>>>(
|
||||
b, c, n, npoints, nsample, points, idx, out);
|
||||
|
||||
CUDA_CHECK_ERRORS();
|
||||
}
|
||||
|
||||
// input: grad_out(b, c, npoints, nsample), idx(b, npoints, nsample)
|
||||
// output: grad_points(b, c, n)
|
||||
__global__ void group_points_grad_kernel(int b, int c, int n, int npoints,
|
||||
int nsample,
|
||||
const float *__restrict__ grad_out,
|
||||
const int *__restrict__ idx,
|
||||
float *__restrict__ grad_points) {
|
||||
int batch_index = blockIdx.x;
|
||||
grad_out += batch_index * npoints * nsample * c;
|
||||
idx += batch_index * npoints * nsample;
|
||||
grad_points += batch_index * n * c;
|
||||
|
||||
const int index = threadIdx.y * blockDim.x + threadIdx.x;
|
||||
const int stride = blockDim.y * blockDim.x;
|
||||
for (int i = index; i < c * npoints; i += stride) {
|
||||
const int l = i / npoints;
|
||||
const int j = i % npoints;
|
||||
for (int k = 0; k < nsample; ++k) {
|
||||
int ii = idx[j * nsample + k];
|
||||
atomicAdd(grad_points + l * n + ii,
|
||||
grad_out[(l * npoints + j) * nsample + k]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints,
|
||||
int nsample, const float *grad_out,
|
||||
const int *idx, float *grad_points) {
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
group_points_grad_kernel<<<b, opt_block_config(npoints, c), 0, stream>>>(
|
||||
b, c, n, npoints, nsample, grad_out, idx, grad_points);
|
||||
|
||||
CUDA_CHECK_ERRORS();
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
#include "interpolate.h"
|
||||
#include "utils.h"
|
||||
|
||||
void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown,
|
||||
const float *known, float *dist2, int *idx);
|
||||
void three_interpolate_kernel_wrapper(int b, int c, int m, int n,
|
||||
const float *points, const int *idx,
|
||||
const float *weight, float *out);
|
||||
void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m,
|
||||
const float *grad_out,
|
||||
const int *idx, const float *weight,
|
||||
float *grad_points);
|
||||
|
||||
std::vector<at::Tensor> three_nn(at::Tensor unknowns, at::Tensor knows) {
|
||||
CHECK_CONTIGUOUS(unknowns);
|
||||
CHECK_CONTIGUOUS(knows);
|
||||
CHECK_IS_FLOAT(unknowns);
|
||||
CHECK_IS_FLOAT(knows);
|
||||
|
||||
if (unknowns.is_cuda()) {
|
||||
CHECK_CUDA(knows);
|
||||
}
|
||||
|
||||
at::Tensor idx =
|
||||
torch::zeros({unknowns.size(0), unknowns.size(1), 3},
|
||||
at::device(unknowns.device()).dtype(at::ScalarType::Int));
|
||||
at::Tensor dist2 =
|
||||
torch::zeros({unknowns.size(0), unknowns.size(1), 3},
|
||||
at::device(unknowns.device()).dtype(at::ScalarType::Float));
|
||||
|
||||
if (unknowns.is_cuda()) {
|
||||
three_nn_kernel_wrapper(unknowns.size(0), unknowns.size(1), knows.size(1),
|
||||
unknowns.data_ptr<float>(), knows.data_ptr<float>(),
|
||||
dist2.data_ptr<float>(), idx.data_ptr<int>());
|
||||
} else {
|
||||
AT_ASSERT(false, "CPU not supported");
|
||||
}
|
||||
|
||||
return {dist2, idx};
|
||||
}
|
||||
|
||||
at::Tensor three_interpolate(at::Tensor points, at::Tensor idx,
|
||||
at::Tensor weight) {
|
||||
CHECK_CONTIGUOUS(points);
|
||||
CHECK_CONTIGUOUS(idx);
|
||||
CHECK_CONTIGUOUS(weight);
|
||||
CHECK_IS_FLOAT(points);
|
||||
CHECK_IS_INT(idx);
|
||||
CHECK_IS_FLOAT(weight);
|
||||
|
||||
if (points.is_cuda()) {
|
||||
CHECK_CUDA(idx);
|
||||
CHECK_CUDA(weight);
|
||||
}
|
||||
|
||||
at::Tensor output =
|
||||
torch::zeros({points.size(0), points.size(1), idx.size(1)},
|
||||
at::device(points.device()).dtype(at::ScalarType::Float));
|
||||
|
||||
if (points.is_cuda()) {
|
||||
three_interpolate_kernel_wrapper(
|
||||
points.size(0), points.size(1), points.size(2), idx.size(1),
|
||||
points.data_ptr<float>(), idx.data_ptr<int>(), weight.data_ptr<float>(),
|
||||
output.data_ptr<float>());
|
||||
} else {
|
||||
AT_ASSERT(false, "CPU not supported");
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx,
|
||||
at::Tensor weight, const int m) {
|
||||
CHECK_CONTIGUOUS(grad_out);
|
||||
CHECK_CONTIGUOUS(idx);
|
||||
CHECK_CONTIGUOUS(weight);
|
||||
CHECK_IS_FLOAT(grad_out);
|
||||
CHECK_IS_INT(idx);
|
||||
CHECK_IS_FLOAT(weight);
|
||||
|
||||
if (grad_out.is_cuda()) {
|
||||
CHECK_CUDA(idx);
|
||||
CHECK_CUDA(weight);
|
||||
}
|
||||
|
||||
at::Tensor output =
|
||||
torch::zeros({grad_out.size(0), grad_out.size(1), m},
|
||||
at::device(grad_out.device()).dtype(at::ScalarType::Float));
|
||||
|
||||
if (grad_out.is_cuda()) {
|
||||
three_interpolate_grad_kernel_wrapper(
|
||||
grad_out.size(0), grad_out.size(1), grad_out.size(2), m,
|
||||
grad_out.data_ptr<float>(), idx.data_ptr<int>(),
|
||||
weight.data_ptr<float>(), output.data_ptr<float>());
|
||||
} else {
|
||||
AT_ASSERT(false, "CPU not supported");
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
#include <math.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
#include "cuda_utils.h"
|
||||
|
||||
// input: unknown(b, n, 3) known(b, m, 3)
|
||||
// output: dist2(b, n, 3), idx(b, n, 3)
|
||||
__global__ void three_nn_kernel(int b, int n, int m,
|
||||
const float *__restrict__ unknown,
|
||||
const float *__restrict__ known,
|
||||
float *__restrict__ dist2,
|
||||
int *__restrict__ idx) {
|
||||
int batch_index = blockIdx.x;
|
||||
unknown += batch_index * n * 3;
|
||||
known += batch_index * m * 3;
|
||||
dist2 += batch_index * n * 3;
|
||||
idx += batch_index * n * 3;
|
||||
|
||||
int index = threadIdx.x;
|
||||
int stride = blockDim.x;
|
||||
for (int j = index; j < n; j += stride) {
|
||||
float ux = unknown[j * 3 + 0];
|
||||
float uy = unknown[j * 3 + 1];
|
||||
float uz = unknown[j * 3 + 2];
|
||||
|
||||
double best1 = 1e40, best2 = 1e40, best3 = 1e40;
|
||||
int besti1 = 0, besti2 = 0, besti3 = 0;
|
||||
for (int k = 0; k < m; ++k) {
|
||||
float x = known[k * 3 + 0];
|
||||
float y = known[k * 3 + 1];
|
||||
float z = known[k * 3 + 2];
|
||||
float d = (ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z);
|
||||
if (d < best1) {
|
||||
best3 = best2;
|
||||
besti3 = besti2;
|
||||
best2 = best1;
|
||||
besti2 = besti1;
|
||||
best1 = d;
|
||||
besti1 = k;
|
||||
} else if (d < best2) {
|
||||
best3 = best2;
|
||||
besti3 = besti2;
|
||||
best2 = d;
|
||||
besti2 = k;
|
||||
} else if (d < best3) {
|
||||
best3 = d;
|
||||
besti3 = k;
|
||||
}
|
||||
}
|
||||
dist2[j * 3 + 0] = best1;
|
||||
dist2[j * 3 + 1] = best2;
|
||||
dist2[j * 3 + 2] = best3;
|
||||
|
||||
idx[j * 3 + 0] = besti1;
|
||||
idx[j * 3 + 1] = besti2;
|
||||
idx[j * 3 + 2] = besti3;
|
||||
}
|
||||
}
|
||||
|
||||
void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown,
|
||||
const float *known, float *dist2, int *idx) {
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
three_nn_kernel<<<b, opt_n_threads(n), 0, stream>>>(b, n, m, unknown, known,
|
||||
dist2, idx);
|
||||
|
||||
CUDA_CHECK_ERRORS();
|
||||
}
|
||||
|
||||
// input: points(b, c, m), idx(b, n, 3), weight(b, n, 3)
|
||||
// output: out(b, c, n)
|
||||
__global__ void three_interpolate_kernel(int b, int c, int m, int n,
|
||||
const float *__restrict__ points,
|
||||
const int *__restrict__ idx,
|
||||
const float *__restrict__ weight,
|
||||
float *__restrict__ out) {
|
||||
int batch_index = blockIdx.x;
|
||||
points += batch_index * m * c;
|
||||
|
||||
idx += batch_index * n * 3;
|
||||
weight += batch_index * n * 3;
|
||||
|
||||
out += batch_index * n * c;
|
||||
|
||||
const int index = threadIdx.y * blockDim.x + threadIdx.x;
|
||||
const int stride = blockDim.y * blockDim.x;
|
||||
for (int i = index; i < c * n; i += stride) {
|
||||
const int l = i / n;
|
||||
const int j = i % n;
|
||||
float w1 = weight[j * 3 + 0];
|
||||
float w2 = weight[j * 3 + 1];
|
||||
float w3 = weight[j * 3 + 2];
|
||||
|
||||
int i1 = idx[j * 3 + 0];
|
||||
int i2 = idx[j * 3 + 1];
|
||||
int i3 = idx[j * 3 + 2];
|
||||
|
||||
out[i] = points[l * m + i1] * w1 + points[l * m + i2] * w2 +
|
||||
points[l * m + i3] * w3;
|
||||
}
|
||||
}
|
||||
|
||||
void three_interpolate_kernel_wrapper(int b, int c, int m, int n,
|
||||
const float *points, const int *idx,
|
||||
const float *weight, float *out) {
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
three_interpolate_kernel<<<b, opt_block_config(n, c), 0, stream>>>(
|
||||
b, c, m, n, points, idx, weight, out);
|
||||
|
||||
CUDA_CHECK_ERRORS();
|
||||
}
|
||||
|
||||
// input: grad_out(b, c, n), idx(b, n, 3), weight(b, n, 3)
|
||||
// output: grad_points(b, c, m)
|
||||
|
||||
__global__ void three_interpolate_grad_kernel(
|
||||
int b, int c, int n, int m, const float *__restrict__ grad_out,
|
||||
const int *__restrict__ idx, const float *__restrict__ weight,
|
||||
float *__restrict__ grad_points) {
|
||||
int batch_index = blockIdx.x;
|
||||
grad_out += batch_index * n * c;
|
||||
idx += batch_index * n * 3;
|
||||
weight += batch_index * n * 3;
|
||||
grad_points += batch_index * m * c;
|
||||
|
||||
const int index = threadIdx.y * blockDim.x + threadIdx.x;
|
||||
const int stride = blockDim.y * blockDim.x;
|
||||
for (int i = index; i < c * n; i += stride) {
|
||||
const int l = i / n;
|
||||
const int j = i % n;
|
||||
float w1 = weight[j * 3 + 0];
|
||||
float w2 = weight[j * 3 + 1];
|
||||
float w3 = weight[j * 3 + 2];
|
||||
|
||||
int i1 = idx[j * 3 + 0];
|
||||
int i2 = idx[j * 3 + 1];
|
||||
int i3 = idx[j * 3 + 2];
|
||||
|
||||
atomicAdd(grad_points + l * m + i1, grad_out[i] * w1);
|
||||
atomicAdd(grad_points + l * m + i2, grad_out[i] * w2);
|
||||
atomicAdd(grad_points + l * m + i3, grad_out[i] * w3);
|
||||
}
|
||||
}
|
||||
|
||||
void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m,
|
||||
const float *grad_out,
|
||||
const int *idx, const float *weight,
|
||||
float *grad_points) {
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
three_interpolate_grad_kernel<<<b, opt_block_config(n, c), 0, stream>>>(
|
||||
b, c, n, m, grad_out, idx, weight, grad_points);
|
||||
|
||||
CUDA_CHECK_ERRORS();
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
#include "sampling.h"
|
||||
#include "utils.h"
|
||||
|
||||
void gather_points_kernel_wrapper(int b, int c, int n, int npoints,
|
||||
const float *points, const int *idx,
|
||||
float *out);
|
||||
void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints,
|
||||
const float *grad_out, const int *idx,
|
||||
float *grad_points);
|
||||
|
||||
void furthest_point_sampling_kernel_wrapper(int b, int n, int m,
|
||||
const float *dataset, float *temp,
|
||||
int *idxs);
|
||||
|
||||
at::Tensor gather_points(at::Tensor points, at::Tensor idx) {
|
||||
CHECK_CONTIGUOUS(points);
|
||||
CHECK_CONTIGUOUS(idx);
|
||||
CHECK_IS_FLOAT(points);
|
||||
CHECK_IS_INT(idx);
|
||||
|
||||
if (points.is_cuda()) {
|
||||
CHECK_CUDA(idx);
|
||||
}
|
||||
|
||||
at::Tensor output =
|
||||
torch::zeros({points.size(0), points.size(1), idx.size(1)},
|
||||
at::device(points.device()).dtype(at::ScalarType::Float));
|
||||
|
||||
if (points.is_cuda()) {
|
||||
gather_points_kernel_wrapper(points.size(0), points.size(1), points.size(2),
|
||||
idx.size(1), points.data_ptr<float>(),
|
||||
idx.data_ptr<int>(), output.data_ptr<float>());
|
||||
} else {
|
||||
AT_ASSERT(false, "CPU not supported");
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx,
|
||||
const int n) {
|
||||
CHECK_CONTIGUOUS(grad_out);
|
||||
CHECK_CONTIGUOUS(idx);
|
||||
CHECK_IS_FLOAT(grad_out);
|
||||
CHECK_IS_INT(idx);
|
||||
|
||||
if (grad_out.is_cuda()) {
|
||||
CHECK_CUDA(idx);
|
||||
}
|
||||
|
||||
at::Tensor output =
|
||||
torch::zeros({grad_out.size(0), grad_out.size(1), n},
|
||||
at::device(grad_out.device()).dtype(at::ScalarType::Float));
|
||||
|
||||
if (grad_out.is_cuda()) {
|
||||
gather_points_grad_kernel_wrapper(grad_out.size(0), grad_out.size(1), n,
|
||||
idx.size(1), grad_out.data_ptr<float>(),
|
||||
idx.data_ptr<int>(),
|
||||
output.data_ptr<float>());
|
||||
} else {
|
||||
AT_ASSERT(false, "CPU not supported");
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples) {
|
||||
CHECK_CONTIGUOUS(points);
|
||||
CHECK_IS_FLOAT(points);
|
||||
|
||||
at::Tensor output =
|
||||
torch::zeros({points.size(0), nsamples},
|
||||
at::device(points.device()).dtype(at::ScalarType::Int));
|
||||
|
||||
at::Tensor tmp =
|
||||
torch::full({points.size(0), points.size(1)}, 1e10,
|
||||
at::device(points.device()).dtype(at::ScalarType::Float));
|
||||
|
||||
if (points.is_cuda()) {
|
||||
furthest_point_sampling_kernel_wrapper(
|
||||
points.size(0), points.size(1), nsamples, points.data_ptr<float>(),
|
||||
tmp.data_ptr<float>(), output.data_ptr<int>());
|
||||
} else {
|
||||
AT_ASSERT(false, "CPU not supported");
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
@@ -0,0 +1,229 @@
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
#include "cuda_utils.h"
|
||||
|
||||
// input: points(b, c, n) idx(b, m)
|
||||
// output: out(b, c, m)
|
||||
__global__ void gather_points_kernel(int b, int c, int n, int m,
|
||||
const float *__restrict__ points,
|
||||
const int *__restrict__ idx,
|
||||
float *__restrict__ out) {
|
||||
for (int i = blockIdx.x; i < b; i += gridDim.x) {
|
||||
for (int l = blockIdx.y; l < c; l += gridDim.y) {
|
||||
for (int j = threadIdx.x; j < m; j += blockDim.x) {
|
||||
int a = idx[i * m + j];
|
||||
out[(i * c + l) * m + j] = points[(i * c + l) * n + a];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void gather_points_kernel_wrapper(int b, int c, int n, int npoints,
|
||||
const float *points, const int *idx,
|
||||
float *out) {
|
||||
gather_points_kernel<<<dim3(b, c, 1), opt_n_threads(npoints), 0,
|
||||
at::cuda::getCurrentCUDAStream()>>>(b, c, n, npoints,
|
||||
points, idx, out);
|
||||
|
||||
CUDA_CHECK_ERRORS();
|
||||
}
|
||||
|
||||
// input: grad_out(b, c, m) idx(b, m)
|
||||
// output: grad_points(b, c, n)
|
||||
__global__ void gather_points_grad_kernel(int b, int c, int n, int m,
|
||||
const float *__restrict__ grad_out,
|
||||
const int *__restrict__ idx,
|
||||
float *__restrict__ grad_points) {
|
||||
for (int i = blockIdx.x; i < b; i += gridDim.x) {
|
||||
for (int l = blockIdx.y; l < c; l += gridDim.y) {
|
||||
for (int j = threadIdx.x; j < m; j += blockDim.x) {
|
||||
int a = idx[i * m + j];
|
||||
atomicAdd(grad_points + (i * c + l) * n + a,
|
||||
grad_out[(i * c + l) * m + j]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints,
|
||||
const float *grad_out, const int *idx,
|
||||
float *grad_points) {
|
||||
gather_points_grad_kernel<<<dim3(b, c, 1), opt_n_threads(npoints), 0,
|
||||
at::cuda::getCurrentCUDAStream()>>>(
|
||||
b, c, n, npoints, grad_out, idx, grad_points);
|
||||
|
||||
CUDA_CHECK_ERRORS();
|
||||
}
|
||||
|
||||
__device__ void __update(float *__restrict__ dists, int *__restrict__ dists_i,
|
||||
int idx1, int idx2) {
|
||||
const float v1 = dists[idx1], v2 = dists[idx2];
|
||||
const int i1 = dists_i[idx1], i2 = dists_i[idx2];
|
||||
dists[idx1] = max(v1, v2);
|
||||
dists_i[idx1] = v2 > v1 ? i2 : i1;
|
||||
}
|
||||
|
||||
// Input dataset: (b, n, 3), tmp: (b, n)
|
||||
// Ouput idxs (b, m)
|
||||
template <unsigned int block_size>
|
||||
__global__ void furthest_point_sampling_kernel(
|
||||
int b, int n, int m, const float *__restrict__ dataset,
|
||||
float *__restrict__ temp, int *__restrict__ idxs) {
|
||||
if (m <= 0) return;
|
||||
__shared__ float dists[block_size];
|
||||
__shared__ int dists_i[block_size];
|
||||
|
||||
int batch_index = blockIdx.x;
|
||||
dataset += batch_index * n * 3;
|
||||
temp += batch_index * n;
|
||||
idxs += batch_index * m;
|
||||
|
||||
int tid = threadIdx.x;
|
||||
const int stride = block_size;
|
||||
|
||||
int old = 0;
|
||||
if (threadIdx.x == 0) idxs[0] = old;
|
||||
|
||||
__syncthreads();
|
||||
for (int j = 1; j < m; j++) {
|
||||
int besti = 0;
|
||||
float best = -1;
|
||||
float x1 = dataset[old * 3 + 0];
|
||||
float y1 = dataset[old * 3 + 1];
|
||||
float z1 = dataset[old * 3 + 2];
|
||||
for (int k = tid; k < n; k += stride) {
|
||||
float x2, y2, z2;
|
||||
x2 = dataset[k * 3 + 0];
|
||||
y2 = dataset[k * 3 + 1];
|
||||
z2 = dataset[k * 3 + 2];
|
||||
float mag = (x2 * x2) + (y2 * y2) + (z2 * z2);
|
||||
if (mag <= 1e-3) continue;
|
||||
|
||||
float d =
|
||||
(x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1);
|
||||
|
||||
float d2 = min(d, temp[k]);
|
||||
temp[k] = d2;
|
||||
besti = d2 > best ? k : besti;
|
||||
best = d2 > best ? d2 : best;
|
||||
}
|
||||
dists[tid] = best;
|
||||
dists_i[tid] = besti;
|
||||
__syncthreads();
|
||||
|
||||
if (block_size >= 512) {
|
||||
if (tid < 256) {
|
||||
__update(dists, dists_i, tid, tid + 256);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (block_size >= 256) {
|
||||
if (tid < 128) {
|
||||
__update(dists, dists_i, tid, tid + 128);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (block_size >= 128) {
|
||||
if (tid < 64) {
|
||||
__update(dists, dists_i, tid, tid + 64);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (block_size >= 64) {
|
||||
if (tid < 32) {
|
||||
__update(dists, dists_i, tid, tid + 32);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (block_size >= 32) {
|
||||
if (tid < 16) {
|
||||
__update(dists, dists_i, tid, tid + 16);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (block_size >= 16) {
|
||||
if (tid < 8) {
|
||||
__update(dists, dists_i, tid, tid + 8);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (block_size >= 8) {
|
||||
if (tid < 4) {
|
||||
__update(dists, dists_i, tid, tid + 4);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (block_size >= 4) {
|
||||
if (tid < 2) {
|
||||
__update(dists, dists_i, tid, tid + 2);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (block_size >= 2) {
|
||||
if (tid < 1) {
|
||||
__update(dists, dists_i, tid, tid + 1);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
old = dists_i[0];
|
||||
if (tid == 0) idxs[j] = old;
|
||||
}
|
||||
}
|
||||
|
||||
void furthest_point_sampling_kernel_wrapper(int b, int n, int m,
|
||||
const float *dataset, float *temp,
|
||||
int *idxs) {
|
||||
unsigned int n_threads = opt_n_threads(n);
|
||||
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
switch (n_threads) {
|
||||
case 512:
|
||||
furthest_point_sampling_kernel<512>
|
||||
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 256:
|
||||
furthest_point_sampling_kernel<256>
|
||||
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 128:
|
||||
furthest_point_sampling_kernel<128>
|
||||
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 64:
|
||||
furthest_point_sampling_kernel<64>
|
||||
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 32:
|
||||
furthest_point_sampling_kernel<32>
|
||||
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 16:
|
||||
furthest_point_sampling_kernel<16>
|
||||
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 8:
|
||||
furthest_point_sampling_kernel<8>
|
||||
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 4:
|
||||
furthest_point_sampling_kernel<4>
|
||||
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 2:
|
||||
furthest_point_sampling_kernel<2>
|
||||
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
case 1:
|
||||
furthest_point_sampling_kernel<1>
|
||||
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
|
||||
break;
|
||||
default:
|
||||
furthest_point_sampling_kernel<512>
|
||||
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
|
||||
}
|
||||
|
||||
CUDA_CHECK_ERRORS();
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
__version__ = "3.0.0"
|
||||
@@ -0,0 +1,209 @@
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from pointnet2_ops import pointnet2_utils
|
||||
|
||||
|
||||
def build_shared_mlp(mlp_spec: List[int], bn: bool = True):
|
||||
layers = []
|
||||
for i in range(1, len(mlp_spec)):
|
||||
layers.append(
|
||||
nn.Conv2d(mlp_spec[i - 1], mlp_spec[i], kernel_size=1, bias=not bn)
|
||||
)
|
||||
if bn:
|
||||
layers.append(nn.BatchNorm2d(mlp_spec[i]))
|
||||
layers.append(nn.ReLU(True))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
|
||||
class _PointnetSAModuleBase(nn.Module):
|
||||
def __init__(self):
|
||||
super(_PointnetSAModuleBase, self).__init__()
|
||||
self.npoint = None
|
||||
self.groupers = None
|
||||
self.mlps = None
|
||||
|
||||
def forward(
|
||||
self, xyz: torch.Tensor, features: Optional[torch.Tensor]
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
r"""
|
||||
Parameters
|
||||
----------
|
||||
xyz : torch.Tensor
|
||||
(B, N, 3) tensor of the xyz coordinates of the features
|
||||
features : torch.Tensor
|
||||
(B, C, N) tensor of the descriptors of the the features
|
||||
|
||||
Returns
|
||||
-------
|
||||
new_xyz : torch.Tensor
|
||||
(B, npoint, 3) tensor of the new features' xyz
|
||||
new_features : torch.Tensor
|
||||
(B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors
|
||||
"""
|
||||
|
||||
new_features_list = []
|
||||
|
||||
xyz_flipped = xyz.transpose(1, 2).contiguous()
|
||||
new_xyz = (
|
||||
pointnet2_utils.gather_operation(
|
||||
xyz_flipped, pointnet2_utils.furthest_point_sample(xyz, self.npoint)
|
||||
)
|
||||
.transpose(1, 2)
|
||||
.contiguous()
|
||||
if self.npoint is not None
|
||||
else None
|
||||
)
|
||||
|
||||
for i in range(len(self.groupers)):
|
||||
new_features = self.groupers[i](
|
||||
xyz, new_xyz, features
|
||||
) # (B, C, npoint, nsample)
|
||||
|
||||
new_features = self.mlps[i](new_features) # (B, mlp[-1], npoint, nsample)
|
||||
new_features = F.max_pool2d(
|
||||
new_features, kernel_size=[1, new_features.size(3)]
|
||||
) # (B, mlp[-1], npoint, 1)
|
||||
new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint)
|
||||
|
||||
new_features_list.append(new_features)
|
||||
|
||||
return new_xyz, torch.cat(new_features_list, dim=1)
|
||||
|
||||
|
||||
class PointnetSAModuleMSG(_PointnetSAModuleBase):
|
||||
r"""Pointnet set abstrction layer with multiscale grouping
|
||||
|
||||
Parameters
|
||||
----------
|
||||
npoint : int
|
||||
Number of features
|
||||
radii : list of float32
|
||||
list of radii to group with
|
||||
nsamples : list of int32
|
||||
Number of samples in each ball query
|
||||
mlps : list of list of int32
|
||||
Spec of the pointnet before the global max_pool for each scale
|
||||
bn : bool
|
||||
Use batchnorm
|
||||
"""
|
||||
|
||||
def __init__(self, npoint, radii, nsamples, mlps, bn=True, use_xyz=True):
|
||||
# type: (PointnetSAModuleMSG, int, List[float], List[int], List[List[int]], bool, bool) -> None
|
||||
super(PointnetSAModuleMSG, self).__init__()
|
||||
|
||||
assert len(radii) == len(nsamples) == len(mlps)
|
||||
|
||||
self.npoint = npoint
|
||||
self.groupers = nn.ModuleList()
|
||||
self.mlps = nn.ModuleList()
|
||||
for i in range(len(radii)):
|
||||
radius = radii[i]
|
||||
nsample = nsamples[i]
|
||||
self.groupers.append(
|
||||
pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz)
|
||||
if npoint is not None
|
||||
else pointnet2_utils.GroupAll(use_xyz)
|
||||
)
|
||||
mlp_spec = mlps[i]
|
||||
if use_xyz:
|
||||
mlp_spec[0] += 3
|
||||
|
||||
self.mlps.append(build_shared_mlp(mlp_spec, bn))
|
||||
|
||||
|
||||
class PointnetSAModule(PointnetSAModuleMSG):
|
||||
r"""Pointnet set abstrction layer
|
||||
|
||||
Parameters
|
||||
----------
|
||||
npoint : int
|
||||
Number of features
|
||||
radius : float
|
||||
Radius of ball
|
||||
nsample : int
|
||||
Number of samples in the ball query
|
||||
mlp : list
|
||||
Spec of the pointnet before the global max_pool
|
||||
bn : bool
|
||||
Use batchnorm
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, mlp, npoint=None, radius=None, nsample=None, bn=True, use_xyz=True
|
||||
):
|
||||
# type: (PointnetSAModule, List[int], int, float, int, bool, bool) -> None
|
||||
super(PointnetSAModule, self).__init__(
|
||||
mlps=[mlp],
|
||||
npoint=npoint,
|
||||
radii=[radius],
|
||||
nsamples=[nsample],
|
||||
bn=bn,
|
||||
use_xyz=use_xyz,
|
||||
)
|
||||
|
||||
|
||||
class PointnetFPModule(nn.Module):
|
||||
r"""Propigates the features of one set to another
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mlp : list
|
||||
Pointnet module parameters
|
||||
bn : bool
|
||||
Use batchnorm
|
||||
"""
|
||||
|
||||
def __init__(self, mlp, bn=True):
|
||||
# type: (PointnetFPModule, List[int], bool) -> None
|
||||
super(PointnetFPModule, self).__init__()
|
||||
self.mlp = build_shared_mlp(mlp, bn=bn)
|
||||
|
||||
def forward(self, unknown, known, unknow_feats, known_feats):
|
||||
# type: (PointnetFPModule, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor) -> torch.Tensor
|
||||
r"""
|
||||
Parameters
|
||||
----------
|
||||
unknown : torch.Tensor
|
||||
(B, n, 3) tensor of the xyz positions of the unknown features
|
||||
known : torch.Tensor
|
||||
(B, m, 3) tensor of the xyz positions of the known features
|
||||
unknow_feats : torch.Tensor
|
||||
(B, C1, n) tensor of the features to be propigated to
|
||||
known_feats : torch.Tensor
|
||||
(B, C2, m) tensor of features to be propigated
|
||||
|
||||
Returns
|
||||
-------
|
||||
new_features : torch.Tensor
|
||||
(B, mlp[-1], n) tensor of the features of the unknown features
|
||||
"""
|
||||
|
||||
if known is not None:
|
||||
dist, idx = pointnet2_utils.three_nn(unknown, known)
|
||||
dist_recip = 1.0 / (dist + 1e-8)
|
||||
norm = torch.sum(dist_recip, dim=2, keepdim=True)
|
||||
weight = dist_recip / norm
|
||||
|
||||
interpolated_feats = pointnet2_utils.three_interpolate(
|
||||
known_feats, idx, weight
|
||||
)
|
||||
else:
|
||||
interpolated_feats = known_feats.expand(
|
||||
*(known_feats.size()[0:2] + [unknown.size(1)])
|
||||
)
|
||||
|
||||
if unknow_feats is not None:
|
||||
new_features = torch.cat(
|
||||
[interpolated_feats, unknow_feats], dim=1
|
||||
) # (B, C2 + C1, n)
|
||||
else:
|
||||
new_features = interpolated_feats
|
||||
|
||||
new_features = new_features.unsqueeze(-1)
|
||||
new_features = self.mlp(new_features)
|
||||
|
||||
return new_features.squeeze(-1)
|
||||
@@ -0,0 +1,379 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import warnings
|
||||
from torch.autograd import Function
|
||||
from typing import *
|
||||
|
||||
try:
|
||||
import pointnet2_ops._ext as _ext
|
||||
except ImportError:
|
||||
from torch.utils.cpp_extension import load
|
||||
import glob
|
||||
import os.path as osp
|
||||
import os
|
||||
|
||||
warnings.warn("Unable to load pointnet2_ops cpp extension. JIT Compiling.")
|
||||
|
||||
_ext_src_root = osp.join(osp.dirname(__file__), "_ext-src")
|
||||
_ext_sources = glob.glob(osp.join(_ext_src_root, "src", "*.cpp")) + glob.glob(
|
||||
osp.join(_ext_src_root, "src", "*.cu")
|
||||
)
|
||||
_ext_headers = glob.glob(osp.join(_ext_src_root, "include", "*"))
|
||||
|
||||
os.environ["TORCH_CUDA_ARCH_LIST"] = "3.7+PTX;5.0;6.0;6.1;6.2;7.0;7.5"
|
||||
_ext = load(
|
||||
"_ext",
|
||||
sources=_ext_sources,
|
||||
extra_include_paths=[osp.join(_ext_src_root, "include")],
|
||||
extra_cflags=["-O3"],
|
||||
extra_cuda_cflags=["-O3", "-Xfatbin", "-compress-all"],
|
||||
with_cuda=True,
|
||||
)
|
||||
|
||||
|
||||
class FurthestPointSampling(Function):
|
||||
@staticmethod
|
||||
def forward(ctx, xyz, npoint):
|
||||
# type: (Any, torch.Tensor, int) -> torch.Tensor
|
||||
r"""
|
||||
Uses iterative furthest point sampling to select a set of npoint features that have the largest
|
||||
minimum distance
|
||||
|
||||
Parameters
|
||||
----------
|
||||
xyz : torch.Tensor
|
||||
(B, N, 3) tensor where N > npoint
|
||||
npoint : int32
|
||||
number of features in the sampled set
|
||||
|
||||
Returns
|
||||
-------
|
||||
torch.Tensor
|
||||
(B, npoint) tensor containing the set
|
||||
"""
|
||||
out = _ext.furthest_point_sampling(xyz, npoint)
|
||||
|
||||
ctx.mark_non_differentiable(out)
|
||||
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_out):
|
||||
return ()
|
||||
|
||||
|
||||
furthest_point_sample = FurthestPointSampling.apply
|
||||
|
||||
|
||||
class GatherOperation(Function):
|
||||
@staticmethod
|
||||
def forward(ctx, features, idx):
|
||||
# type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor
|
||||
r"""
|
||||
|
||||
Parameters
|
||||
----------
|
||||
features : torch.Tensor
|
||||
(B, C, N) tensor
|
||||
|
||||
idx : torch.Tensor
|
||||
(B, npoint) tensor of the features to gather
|
||||
|
||||
Returns
|
||||
-------
|
||||
torch.Tensor
|
||||
(B, C, npoint) tensor
|
||||
"""
|
||||
|
||||
ctx.save_for_backward(idx, features)
|
||||
|
||||
return _ext.gather_points(features, idx)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_out):
|
||||
idx, features = ctx.saved_tensors
|
||||
N = features.size(2)
|
||||
|
||||
grad_features = _ext.gather_points_grad(grad_out.contiguous(), idx, N)
|
||||
return grad_features, None
|
||||
|
||||
|
||||
gather_operation = GatherOperation.apply
|
||||
|
||||
|
||||
class ThreeNN(Function):
|
||||
@staticmethod
|
||||
def forward(ctx, unknown, known):
|
||||
# type: (Any, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]
|
||||
r"""
|
||||
Find the three nearest neighbors of unknown in known
|
||||
Parameters
|
||||
----------
|
||||
unknown : torch.Tensor
|
||||
(B, n, 3) tensor of known features
|
||||
known : torch.Tensor
|
||||
(B, m, 3) tensor of unknown features
|
||||
|
||||
Returns
|
||||
-------
|
||||
dist : torch.Tensor
|
||||
(B, n, 3) l2 distance to the three nearest neighbors
|
||||
idx : torch.Tensor
|
||||
(B, n, 3) index of 3 nearest neighbors
|
||||
"""
|
||||
dist2, idx = _ext.three_nn(unknown, known)
|
||||
dist = torch.sqrt(dist2)
|
||||
|
||||
ctx.mark_non_differentiable(dist, idx)
|
||||
|
||||
return dist, idx
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_dist, grad_idx):
|
||||
return ()
|
||||
|
||||
|
||||
three_nn = ThreeNN.apply
|
||||
|
||||
|
||||
class ThreeInterpolate(Function):
|
||||
@staticmethod
|
||||
def forward(ctx, features, idx, weight):
|
||||
# type(Any, torch.Tensor, torch.Tensor, torch.Tensor) -> Torch.Tensor
|
||||
r"""
|
||||
Performs weight linear interpolation on 3 features
|
||||
Parameters
|
||||
----------
|
||||
features : torch.Tensor
|
||||
(B, c, m) Features descriptors to be interpolated from
|
||||
idx : torch.Tensor
|
||||
(B, n, 3) three nearest neighbors of the target features in features
|
||||
weight : torch.Tensor
|
||||
(B, n, 3) weights
|
||||
|
||||
Returns
|
||||
-------
|
||||
torch.Tensor
|
||||
(B, c, n) tensor of the interpolated features
|
||||
"""
|
||||
ctx.save_for_backward(idx, weight, features)
|
||||
|
||||
return _ext.three_interpolate(features, idx, weight)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_out):
|
||||
# type: (Any, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
|
||||
r"""
|
||||
Parameters
|
||||
----------
|
||||
grad_out : torch.Tensor
|
||||
(B, c, n) tensor with gradients of ouputs
|
||||
|
||||
Returns
|
||||
-------
|
||||
grad_features : torch.Tensor
|
||||
(B, c, m) tensor with gradients of features
|
||||
|
||||
None
|
||||
|
||||
None
|
||||
"""
|
||||
idx, weight, features = ctx.saved_tensors
|
||||
m = features.size(2)
|
||||
|
||||
grad_features = _ext.three_interpolate_grad(
|
||||
grad_out.contiguous(), idx, weight, m
|
||||
)
|
||||
|
||||
return grad_features, torch.zeros_like(idx), torch.zeros_like(weight)
|
||||
|
||||
|
||||
three_interpolate = ThreeInterpolate.apply
|
||||
|
||||
|
||||
class GroupingOperation(Function):
|
||||
@staticmethod
|
||||
def forward(ctx, features, idx):
|
||||
# type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor
|
||||
r"""
|
||||
|
||||
Parameters
|
||||
----------
|
||||
features : torch.Tensor
|
||||
(B, C, N) tensor of features to group
|
||||
idx : torch.Tensor
|
||||
(B, npoint, nsample) tensor containing the indicies of features to group with
|
||||
|
||||
Returns
|
||||
-------
|
||||
torch.Tensor
|
||||
(B, C, npoint, nsample) tensor
|
||||
"""
|
||||
ctx.save_for_backward(idx, features)
|
||||
|
||||
return _ext.group_points(features, idx)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_out):
|
||||
# type: (Any, torch.tensor) -> Tuple[torch.Tensor, torch.Tensor]
|
||||
r"""
|
||||
|
||||
Parameters
|
||||
----------
|
||||
grad_out : torch.Tensor
|
||||
(B, C, npoint, nsample) tensor of the gradients of the output from forward
|
||||
|
||||
Returns
|
||||
-------
|
||||
torch.Tensor
|
||||
(B, C, N) gradient of the features
|
||||
None
|
||||
"""
|
||||
idx, features = ctx.saved_tensors
|
||||
N = features.size(2)
|
||||
|
||||
grad_features = _ext.group_points_grad(grad_out.contiguous(), idx, N)
|
||||
|
||||
return grad_features, torch.zeros_like(idx)
|
||||
|
||||
|
||||
grouping_operation = GroupingOperation.apply
|
||||
|
||||
|
||||
class BallQuery(Function):
|
||||
@staticmethod
|
||||
def forward(ctx, radius, nsample, xyz, new_xyz):
|
||||
# type: (Any, float, int, torch.Tensor, torch.Tensor) -> torch.Tensor
|
||||
r"""
|
||||
|
||||
Parameters
|
||||
----------
|
||||
radius : float
|
||||
radius of the balls
|
||||
nsample : int
|
||||
maximum number of features in the balls
|
||||
xyz : torch.Tensor
|
||||
(B, N, 3) xyz coordinates of the features
|
||||
new_xyz : torch.Tensor
|
||||
(B, npoint, 3) centers of the ball query
|
||||
|
||||
Returns
|
||||
-------
|
||||
torch.Tensor
|
||||
(B, npoint, nsample) tensor with the indicies of the features that form the query balls
|
||||
"""
|
||||
output = _ext.ball_query(new_xyz, xyz, radius, nsample)
|
||||
|
||||
ctx.mark_non_differentiable(output)
|
||||
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_out):
|
||||
return ()
|
||||
|
||||
|
||||
ball_query = BallQuery.apply
|
||||
|
||||
|
||||
class QueryAndGroup(nn.Module):
|
||||
r"""
|
||||
Groups with a ball query of radius
|
||||
|
||||
Parameters
|
||||
---------
|
||||
radius : float32
|
||||
Radius of ball
|
||||
nsample : int32
|
||||
Maximum number of features to gather in the ball
|
||||
"""
|
||||
|
||||
def __init__(self, radius, nsample, use_xyz=True):
|
||||
# type: (QueryAndGroup, float, int, bool) -> None
|
||||
super(QueryAndGroup, self).__init__()
|
||||
self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz
|
||||
|
||||
def forward(self, xyz, new_xyz, features=None):
|
||||
# type: (QueryAndGroup, torch.Tensor. torch.Tensor, torch.Tensor) -> Tuple[Torch.Tensor]
|
||||
r"""
|
||||
Parameters
|
||||
----------
|
||||
xyz : torch.Tensor
|
||||
xyz coordinates of the features (B, N, 3)
|
||||
new_xyz : torch.Tensor
|
||||
centriods (B, npoint, 3)
|
||||
features : torch.Tensor
|
||||
Descriptors of the features (B, C, N)
|
||||
|
||||
Returns
|
||||
-------
|
||||
new_features : torch.Tensor
|
||||
(B, 3 + C, npoint, nsample) tensor
|
||||
"""
|
||||
|
||||
idx = ball_query(self.radius, self.nsample, xyz, new_xyz)
|
||||
xyz_trans = xyz.transpose(1, 2).contiguous()
|
||||
grouped_xyz = grouping_operation(xyz_trans, idx) # (B, 3, npoint, nsample)
|
||||
grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1)
|
||||
|
||||
if features is not None:
|
||||
grouped_features = grouping_operation(features, idx)
|
||||
if self.use_xyz:
|
||||
new_features = torch.cat(
|
||||
[grouped_xyz, grouped_features], dim=1
|
||||
) # (B, C + 3, npoint, nsample)
|
||||
else:
|
||||
new_features = grouped_features
|
||||
else:
|
||||
assert (
|
||||
self.use_xyz
|
||||
), "Cannot have not features and not use xyz as a feature!"
|
||||
new_features = grouped_xyz
|
||||
|
||||
return new_features
|
||||
|
||||
|
||||
class GroupAll(nn.Module):
|
||||
r"""
|
||||
Groups all features
|
||||
|
||||
Parameters
|
||||
---------
|
||||
"""
|
||||
|
||||
def __init__(self, use_xyz=True):
|
||||
# type: (GroupAll, bool) -> None
|
||||
super(GroupAll, self).__init__()
|
||||
self.use_xyz = use_xyz
|
||||
|
||||
def forward(self, xyz, new_xyz, features=None):
|
||||
# type: (GroupAll, torch.Tensor, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor]
|
||||
r"""
|
||||
Parameters
|
||||
----------
|
||||
xyz : torch.Tensor
|
||||
xyz coordinates of the features (B, N, 3)
|
||||
new_xyz : torch.Tensor
|
||||
Ignored
|
||||
features : torch.Tensor
|
||||
Descriptors of the features (B, C, N)
|
||||
|
||||
Returns
|
||||
-------
|
||||
new_features : torch.Tensor
|
||||
(B, C + 3, 1, N) tensor
|
||||
"""
|
||||
|
||||
grouped_xyz = xyz.transpose(1, 2).unsqueeze(2)
|
||||
if features is not None:
|
||||
grouped_features = features.unsqueeze(2)
|
||||
if self.use_xyz:
|
||||
new_features = torch.cat(
|
||||
[grouped_xyz, grouped_features], dim=1
|
||||
) # (B, 3 + C, 1, N)
|
||||
else:
|
||||
new_features = grouped_features
|
||||
else:
|
||||
new_features = grouped_xyz
|
||||
|
||||
return new_features
|
||||
@@ -0,0 +1,41 @@
|
||||
import glob
|
||||
import os
|
||||
import os.path as osp
|
||||
|
||||
from setuptools import find_packages, setup
|
||||
import torch
|
||||
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||||
|
||||
this_dir = osp.dirname(osp.abspath(__file__))
|
||||
_ext_src_root = osp.join("pointnet2_ops", "_ext-src")
|
||||
_ext_sources = glob.glob(osp.join(_ext_src_root, "src", "*.cpp")) + glob.glob(
|
||||
osp.join(_ext_src_root, "src", "*.cu")
|
||||
)
|
||||
_ext_headers = glob.glob(osp.join(_ext_src_root, "include", "*"))
|
||||
|
||||
requirements = []
|
||||
|
||||
exec(open(osp.join("pointnet2_ops", "_version.py")).read())
|
||||
|
||||
os.environ["TORCH_CUDA_ARCH_LIST"] = ".".join(map(str, torch.cuda.get_device_capability()))
|
||||
# os.environ["TORCH_CUDA_ARCH_LIST"] = "5.0;6.0;6.1;6.2;7.0;7.5;8.0;8.6"
|
||||
setup(
|
||||
name="pointnet2_ops",
|
||||
version=__version__,
|
||||
author="Erik Wijmans",
|
||||
packages=find_packages(),
|
||||
install_requires=requirements,
|
||||
ext_modules=[
|
||||
CUDAExtension(
|
||||
name="pointnet2_ops._ext",
|
||||
sources=_ext_sources,
|
||||
extra_compile_args={
|
||||
"cxx": ["-O3"],
|
||||
"nvcc": ["-O3", "-Xfatbin", "-compress-all"],
|
||||
},
|
||||
include_dirs=[osp.join(this_dir, _ext_src_root, "include")],
|
||||
)
|
||||
],
|
||||
cmdclass={"build_ext": BuildExtension},
|
||||
include_package_data=True,
|
||||
)
|
||||
@@ -0,0 +1,17 @@
|
||||
/*
|
||||
* Copyright (C) 2023, Inria
|
||||
* GRAPHDECO research group, https://team.inria.fr/graphdeco
|
||||
* All rights reserved.
|
||||
*
|
||||
* This software is free for non-commercial, research and evaluation use
|
||||
* under the terms of the LICENSE.md file.
|
||||
*
|
||||
* For inquiries contact george.drettakis@inria.fr
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include "spatial.h"
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("distCUDA2", &distCUDA2);
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
#
|
||||
# Copyright (C) 2023, Inria
|
||||
# GRAPHDECO research group, https://team.inria.fr/graphdeco
|
||||
# All rights reserved.
|
||||
#
|
||||
# This software is free for non-commercial, research and evaluation use
|
||||
# under the terms of the LICENSE.md file.
|
||||
#
|
||||
# For inquiries contact george.drettakis@inria.fr
|
||||
#
|
||||
|
||||
from setuptools import setup
|
||||
from torch.utils.cpp_extension import CUDAExtension, BuildExtension
|
||||
import os
|
||||
|
||||
cxx_compiler_flags = []
|
||||
|
||||
if os.name == 'nt':
|
||||
cxx_compiler_flags.append("/wd4624")
|
||||
|
||||
setup(
|
||||
name="simple_knn",
|
||||
ext_modules=[
|
||||
CUDAExtension(
|
||||
name="simple_knn._C",
|
||||
sources=[
|
||||
"spatial.cu",
|
||||
"simple_knn.cu",
|
||||
"ext.cpp"],
|
||||
extra_compile_args={"nvcc": [], "cxx": cxx_compiler_flags})
|
||||
],
|
||||
cmdclass={
|
||||
'build_ext': BuildExtension
|
||||
}
|
||||
)
|
||||
@@ -0,0 +1,222 @@
|
||||
/*
|
||||
* Copyright (C) 2023, Inria
|
||||
* GRAPHDECO research group, https://team.inria.fr/graphdeco
|
||||
* All rights reserved.
|
||||
*
|
||||
* This software is free for non-commercial, research and evaluation use
|
||||
* under the terms of the LICENSE.md file.
|
||||
*
|
||||
* For inquiries contact george.drettakis@inria.fr
|
||||
*/
|
||||
|
||||
#define BOX_SIZE 1024
|
||||
|
||||
#include "cuda_runtime.h"
|
||||
#include "device_launch_parameters.h"
|
||||
#include "simple_knn.h"
|
||||
#include <cub/cub.cuh>
|
||||
#include <cub/device/device_radix_sort.cuh>
|
||||
#include <vector>
|
||||
#include <cuda_runtime_api.h>
|
||||
#include <thrust/device_vector.h>
|
||||
#include <thrust/sequence.h>
|
||||
#define __CUDACC__
|
||||
#include <cooperative_groups.h>
|
||||
#include <cooperative_groups/reduce.h>
|
||||
#include <cfloat>
|
||||
|
||||
namespace cg = cooperative_groups;
|
||||
|
||||
struct CustomMin
|
||||
{
|
||||
__device__ __forceinline__
|
||||
float3 operator()(const float3& a, const float3& b) const {
|
||||
return { min(a.x, b.x), min(a.y, b.y), min(a.z, b.z) };
|
||||
}
|
||||
};
|
||||
|
||||
struct CustomMax
|
||||
{
|
||||
__device__ __forceinline__
|
||||
float3 operator()(const float3& a, const float3& b) const {
|
||||
return { max(a.x, b.x), max(a.y, b.y), max(a.z, b.z) };
|
||||
}
|
||||
};
|
||||
|
||||
__host__ __device__ uint32_t prepMorton(uint32_t x)
|
||||
{
|
||||
x = (x | (x << 16)) & 0x030000FF;
|
||||
x = (x | (x << 8)) & 0x0300F00F;
|
||||
x = (x | (x << 4)) & 0x030C30C3;
|
||||
x = (x | (x << 2)) & 0x09249249;
|
||||
return x;
|
||||
}
|
||||
|
||||
__host__ __device__ uint32_t coord2Morton(float3 coord, float3 minn, float3 maxx)
|
||||
{
|
||||
uint32_t x = prepMorton(((coord.x - minn.x) / (maxx.x - minn.x)) * ((1 << 10) - 1));
|
||||
uint32_t y = prepMorton(((coord.y - minn.y) / (maxx.y - minn.y)) * ((1 << 10) - 1));
|
||||
uint32_t z = prepMorton(((coord.z - minn.z) / (maxx.z - minn.z)) * ((1 << 10) - 1));
|
||||
|
||||
return x | (y << 1) | (z << 2);
|
||||
}
|
||||
|
||||
__global__ void coord2Morton(int P, const float3* points, float3 minn, float3 maxx, uint32_t* codes)
|
||||
{
|
||||
auto idx = cg::this_grid().thread_rank();
|
||||
if (idx >= P)
|
||||
return;
|
||||
|
||||
codes[idx] = coord2Morton(points[idx], minn, maxx);
|
||||
}
|
||||
|
||||
struct MinMax
|
||||
{
|
||||
float3 minn;
|
||||
float3 maxx;
|
||||
};
|
||||
|
||||
__global__ void boxMinMax(uint32_t P, float3* points, uint32_t* indices, MinMax* boxes)
|
||||
{
|
||||
auto idx = cg::this_grid().thread_rank();
|
||||
|
||||
MinMax me;
|
||||
if (idx < P)
|
||||
{
|
||||
me.minn = points[indices[idx]];
|
||||
me.maxx = points[indices[idx]];
|
||||
}
|
||||
else
|
||||
{
|
||||
me.minn = { FLT_MAX, FLT_MAX, FLT_MAX };
|
||||
me.maxx = { -FLT_MAX,-FLT_MAX,-FLT_MAX };
|
||||
}
|
||||
|
||||
__shared__ MinMax redResult[BOX_SIZE];
|
||||
|
||||
for (int off = BOX_SIZE / 2; off >= 1; off /= 2)
|
||||
{
|
||||
if (threadIdx.x < 2 * off)
|
||||
redResult[threadIdx.x] = me;
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x < off)
|
||||
{
|
||||
MinMax other = redResult[threadIdx.x + off];
|
||||
me.minn.x = min(me.minn.x, other.minn.x);
|
||||
me.minn.y = min(me.minn.y, other.minn.y);
|
||||
me.minn.z = min(me.minn.z, other.minn.z);
|
||||
me.maxx.x = max(me.maxx.x, other.maxx.x);
|
||||
me.maxx.y = max(me.maxx.y, other.maxx.y);
|
||||
me.maxx.z = max(me.maxx.z, other.maxx.z);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (threadIdx.x == 0)
|
||||
boxes[blockIdx.x] = me;
|
||||
}
|
||||
|
||||
__device__ __host__ float distBoxPoint(const MinMax& box, const float3& p)
|
||||
{
|
||||
float3 diff = { 0, 0, 0 };
|
||||
if (p.x < box.minn.x || p.x > box.maxx.x)
|
||||
diff.x = min(abs(p.x - box.minn.x), abs(p.x - box.maxx.x));
|
||||
if (p.y < box.minn.y || p.y > box.maxx.y)
|
||||
diff.y = min(abs(p.y - box.minn.y), abs(p.y - box.maxx.y));
|
||||
if (p.z < box.minn.z || p.z > box.maxx.z)
|
||||
diff.z = min(abs(p.z - box.minn.z), abs(p.z - box.maxx.z));
|
||||
return diff.x * diff.x + diff.y * diff.y + diff.z * diff.z;
|
||||
}
|
||||
|
||||
template<int K>
|
||||
__device__ void updateKBest(const float3& ref, const float3& point, float* knn)
|
||||
{
|
||||
float3 d = { point.x - ref.x, point.y - ref.y, point.z - ref.z };
|
||||
float dist = d.x * d.x + d.y * d.y + d.z * d.z;
|
||||
for (int j = 0; j < K; j++)
|
||||
{
|
||||
if (knn[j] > dist)
|
||||
{
|
||||
float t = knn[j];
|
||||
knn[j] = dist;
|
||||
dist = t;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void boxMeanDist(uint32_t P, float3* points, uint32_t* indices, MinMax* boxes, float* dists)
|
||||
{
|
||||
int idx = cg::this_grid().thread_rank();
|
||||
if (idx >= P)
|
||||
return;
|
||||
|
||||
float3 point = points[indices[idx]];
|
||||
float best[3] = { FLT_MAX, FLT_MAX, FLT_MAX };
|
||||
|
||||
for (int i = max(0, idx - 3); i <= min(P - 1, idx + 3); i++)
|
||||
{
|
||||
if (i == idx)
|
||||
continue;
|
||||
updateKBest<3>(point, points[indices[i]], best);
|
||||
}
|
||||
|
||||
float reject = best[2];
|
||||
best[0] = FLT_MAX;
|
||||
best[1] = FLT_MAX;
|
||||
best[2] = FLT_MAX;
|
||||
|
||||
for (int b = 0; b < (P + BOX_SIZE - 1) / BOX_SIZE; b++)
|
||||
{
|
||||
MinMax box = boxes[b];
|
||||
float dist = distBoxPoint(box, point);
|
||||
if (dist > reject || dist > best[2])
|
||||
continue;
|
||||
|
||||
for (int i = b * BOX_SIZE; i < min(P, (b + 1) * BOX_SIZE); i++)
|
||||
{
|
||||
if (i == idx)
|
||||
continue;
|
||||
updateKBest<3>(point, points[indices[i]], best);
|
||||
}
|
||||
}
|
||||
dists[indices[idx]] = (best[0] + best[1] + best[2]) / 3.0f;
|
||||
}
|
||||
|
||||
void SimpleKNN::knn(int P, float3* points, float* meanDists)
|
||||
{
|
||||
float3* result;
|
||||
cudaMalloc(&result, sizeof(float3));
|
||||
size_t temp_storage_bytes;
|
||||
|
||||
float3 init = { 0, 0, 0 }, minn, maxx;
|
||||
|
||||
cub::DeviceReduce::Reduce(nullptr, temp_storage_bytes, points, result, P, CustomMin(), init);
|
||||
thrust::device_vector<char> temp_storage(temp_storage_bytes);
|
||||
|
||||
cub::DeviceReduce::Reduce(temp_storage.data().get(), temp_storage_bytes, points, result, P, CustomMin(), init);
|
||||
cudaMemcpy(&minn, result, sizeof(float3), cudaMemcpyDeviceToHost);
|
||||
|
||||
cub::DeviceReduce::Reduce(temp_storage.data().get(), temp_storage_bytes, points, result, P, CustomMax(), init);
|
||||
cudaMemcpy(&maxx, result, sizeof(float3), cudaMemcpyDeviceToHost);
|
||||
|
||||
thrust::device_vector<uint32_t> morton(P);
|
||||
thrust::device_vector<uint32_t> morton_sorted(P);
|
||||
coord2Morton << <(P + 255) / 256, 256 >> > (P, points, minn, maxx, morton.data().get());
|
||||
|
||||
thrust::device_vector<uint32_t> indices(P);
|
||||
thrust::sequence(indices.begin(), indices.end());
|
||||
thrust::device_vector<uint32_t> indices_sorted(P);
|
||||
|
||||
cub::DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes, morton.data().get(), morton_sorted.data().get(), indices.data().get(), indices_sorted.data().get(), P);
|
||||
temp_storage.resize(temp_storage_bytes);
|
||||
|
||||
cub::DeviceRadixSort::SortPairs(temp_storage.data().get(), temp_storage_bytes, morton.data().get(), morton_sorted.data().get(), indices.data().get(), indices_sorted.data().get(), P);
|
||||
|
||||
uint32_t num_boxes = (P + BOX_SIZE - 1) / BOX_SIZE;
|
||||
thrust::device_vector<MinMax> boxes(num_boxes);
|
||||
boxMinMax << <num_boxes, BOX_SIZE >> > (P, points, indices_sorted.data().get(), boxes.data().get());
|
||||
boxMeanDist << <num_boxes, BOX_SIZE >> > (P, points, indices_sorted.data().get(), boxes.data().get(), meanDists);
|
||||
|
||||
cudaFree(result);
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
/*
|
||||
* Copyright (C) 2023, Inria
|
||||
* GRAPHDECO research group, https://team.inria.fr/graphdeco
|
||||
* All rights reserved.
|
||||
*
|
||||
* This software is free for non-commercial, research and evaluation use
|
||||
* under the terms of the LICENSE.md file.
|
||||
*
|
||||
* For inquiries contact george.drettakis@inria.fr
|
||||
*/
|
||||
|
||||
#ifndef SIMPLEKNN_H_INCLUDED
|
||||
#define SIMPLEKNN_H_INCLUDED
|
||||
|
||||
class SimpleKNN
|
||||
{
|
||||
public:
|
||||
static void knn(int P, float3* points, float* meanDists);
|
||||
};
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,26 @@
|
||||
/*
|
||||
* Copyright (C) 2023, Inria
|
||||
* GRAPHDECO research group, https://team.inria.fr/graphdeco
|
||||
* All rights reserved.
|
||||
*
|
||||
* This software is free for non-commercial, research and evaluation use
|
||||
* under the terms of the LICENSE.md file.
|
||||
*
|
||||
* For inquiries contact george.drettakis@inria.fr
|
||||
*/
|
||||
|
||||
#include "spatial.h"
|
||||
#include "simple_knn.h"
|
||||
|
||||
torch::Tensor
|
||||
distCUDA2(const torch::Tensor& points)
|
||||
{
|
||||
const int P = points.size(0);
|
||||
|
||||
auto float_opts = points.options().dtype(torch::kFloat32);
|
||||
torch::Tensor means = torch::full({P}, 0.0, float_opts);
|
||||
|
||||
SimpleKNN::knn(P, (float3*)points.contiguous().data<float>(), means.contiguous().data<float>());
|
||||
|
||||
return means;
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
/*
|
||||
* Copyright (C) 2023, Inria
|
||||
* GRAPHDECO research group, https://team.inria.fr/graphdeco
|
||||
* All rights reserved.
|
||||
*
|
||||
* This software is free for non-commercial, research and evaluation use
|
||||
* under the terms of the LICENSE.md file.
|
||||
*
|
||||
* For inquiries contact george.drettakis@inria.fr
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
|
||||
torch::Tensor distCUDA2(const torch::Tensor& points);
|
||||
BIN
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -1,6 +1,6 @@
|
||||
# local
|
||||
./Gen_3D_Modules/TriplaneGaussian/models/snowflake/pointnet2_ops_lib
|
||||
./Libs/simple-knn
|
||||
./_Pre_Builds/_Libs/pointnet2_ops
|
||||
./_Pre_Builds/_Libs/simple-knn
|
||||
|
||||
# remote
|
||||
git+https://github.com/ashawkey/diff-gaussian-rasterization
|
||||
|
||||
Reference in New Issue
Block a user