Skip to content

Commit 6840125

Browse files
lgraylgray
committed
results of initial developments
1 parent 3fbfd1a commit 6840125

32 files changed

+2557
-0
lines changed

‎CMakeLists.txt

+77
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
cmake_minimum_required(VERSION 3.0)
2+
project(torchcmspepr)
3+
set(CMAKE_CXX_STANDARD 14)
4+
set(TORCHCMSPEPR_VERSION 0.0.1)
5+
6+
option(WITH_CUDA "Enable CUDA support" OFF)
7+
8+
if(WITH_CUDA)
9+
enable_language(CUDA)
10+
add_definitions(-D__CUDA_NO_HALF_OPERATORS__)
11+
add_definitions(-DWITH_CUDA)
12+
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -arch=sm_35 --expt-relaxed-constexpr")
13+
endif()
14+
15+
find_package(Python3 COMPONENTS Development)
16+
find_package(Torch REQUIRED)
17+
18+
include_directories(csrc/)
19+
file(GLOB HEADERS csrc/cmspepr.h csrc/helpers.h)
20+
file(GLOB OPERATOR_SOURCES csrc/cpu/*.h csrc/cpu/*.cpp csrc/*.cpp)
21+
if(WITH_CUDA)
22+
file(GLOB OPERATOR_SOURCES ${OPERATOR_SOURCES} csrc/cuda/*.h csrc/cuda/*.cu)
23+
endif()
24+
25+
add_library(${PROJECT_NAME} SHARED ${OPERATOR_SOURCES})
26+
target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES} Python3::Python)
27+
set_target_properties(${PROJECT_NAME} PROPERTIES EXPORT_NAME TorchCMSPepr)
28+
29+
target_include_directories(${PROJECT_NAME} INTERFACE
30+
$<BUILD_INTERFACE:${HEADERS}>
31+
$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>)
32+
33+
include(GNUInstallDirs)
34+
include(CMakePackageConfigHelpers)
35+
36+
set(TORCHCMSPEPR_CMAKECONFIG_INSTALL_DIR "share/cmake/TorchCMSPepr" CACHE STRING "install path for TorchCMSPeprConfig.cmake")
37+
38+
configure_package_config_file(cmake/TorchCMSPeprConfig.cmake.in
39+
"${CMAKE_CURRENT_BINARY_DIR}/TorchCMSPeprConfig.cmake"
40+
INSTALL_DESTINATION ${TORCHCMSPEPR_CMAKECONFIG_INSTALL_DIR})
41+
42+
write_basic_package_version_file(${CMAKE_CURRENT_BINARY_DIR}/TorchCMSPeprConfigVersion.cmake
43+
VERSION ${TORCHCMSPEPR_VERSION}
44+
COMPATIBILITY AnyNewerVersion)
45+
46+
install(FILES ${CMAKE_CURRENT_BINARY_DIR}/TorchCMSPeprConfig.cmake
47+
${CMAKE_CURRENT_BINARY_DIR}/TorchCMSPeprConfigVersion.cmake
48+
DESTINATION ${TORCHCMSPEPR_CMAKECONFIG_INSTALL_DIR})
49+
50+
install(TARGETS ${PROJECT_NAME}
51+
EXPORT TorchCMSPeprTargets
52+
LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
53+
)
54+
55+
install(EXPORT TorchCMSPeprTargets
56+
NAMESPACE TorchCMSPepr::
57+
DESTINATION ${TORCHCMSPEPR_CMAKECONFIG_INSTALL_DIR})
58+
59+
install(FILES ${HEADERS} DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME})
60+
install(FILES
61+
csrc/cpu/accumulate_knn_cpu.h
62+
csrc/cpu/accumulate_knn_grad_cpu.h
63+
csrc/cpu/select_knn_cpu.h
64+
DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}/cpu)
65+
if(WITH_CUDA)
66+
install(FILES
67+
csrc/cuda/accumulate_knn_cuda.h
68+
csrc/cuda/accumulate_knn_grad_cuda.h
69+
csrc/cuda/select_knn_cuda.h
70+
csrc/cuda/select_knn_grad_cuda.h
71+
DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}/cuda)
72+
endif()
73+
74+
if(WITH_CUDA)
75+
set_property(TARGET torch_cuda PROPERTY INTERFACE_COMPILE_OPTIONS "")
76+
set_property(TARGET torch_cpu PROPERTY INTERFACE_COMPILE_OPTIONS "")
77+
endif()

‎cmake/TorchCMSPeprConfig.cmake.in

+38
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
# TorchClusterConfig.cmake
2+
# --------------------
3+
#
4+
# Exported targets:: Cluster
5+
#
6+
7+
@PACKAGE_INIT@
8+
9+
set(PN TorchCMSPepr)
10+
set(${PN}_INCLUDE_DIR "${PACKAGE_PREFIX_DIR}/@CMAKE_INSTALL_INCLUDEDIR@")
11+
set(${PN}_LIBRARY "")
12+
set(${PN}_DEFINITIONS USING_${PN})
13+
14+
check_required_components(${PN})
15+
16+
17+
if(NOT (CMAKE_VERSION VERSION_LESS 3.0))
18+
#-----------------------------------------------------------------------------
19+
# Don't include targets if this file is being picked up by another
20+
# project which has already built this as a subproject
21+
#-----------------------------------------------------------------------------
22+
if(NOT TARGET ${PN}::TorchCMSPepr)
23+
include("${CMAKE_CURRENT_LIST_DIR}/${PN}Targets.cmake")
24+
25+
if(NOT TARGET torch_library)
26+
find_package(Torch REQUIRED)
27+
endif()
28+
if(NOT TARGET Python3::Python)
29+
find_package(Python3 COMPONENTS Development)
30+
endif()
31+
target_link_libraries(TorchCMSPepr::TorchCMSPepr INTERFACE ${TORCH_LIBRARIES} Python3::Python)
32+
33+
if(@WITH_CUDA@)
34+
target_compile_definitions(TorchCMSPepr::TorchCMSPepr INTERFACE WITH_CUDA)
35+
endif()
36+
37+
endif()
38+
endif()

‎csrc/accumulate_knn.cpp

+45
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
#include <Python.h>
2+
#include <torch/script.h>
3+
4+
#include "cpu/accumulate_knn_cpu.h"
5+
6+
#ifdef WITH_CUDA
7+
#include "cuda/accumulate_knn_cuda.h"
8+
#endif
9+
10+
11+
#ifdef _WIN32
12+
#ifdef WITH_CUDA
13+
PyMODINIT_FUNC PyInit__accumulate_knn_cuda(void) { return NULL; }
14+
#else
15+
PyMODINIT_FUNC PyInit__accumulate_knn_cpu(void) { return NULL; }
16+
#endif
17+
#endif
18+
19+
std::tuple<torch::Tensor, torch::Tensor>
20+
accumulate_knn(torch::Tensor distances,
21+
torch::Tensor features,
22+
torch::Tensor indices,
23+
int64_t n_moments,
24+
bool mean_and_max) {
25+
if (distances.device().is_cuda()) {
26+
#ifdef WITH_CUDA
27+
return accumulate_knn_cuda(distances,
28+
features,
29+
indices,
30+
n_moments,
31+
mean_and_max);
32+
#else
33+
AT_ERROR("Not compiled with CUDA support");
34+
#endif
35+
} else {
36+
return accumulate_knn_cpu(distances,
37+
features,
38+
indices,
39+
n_moments,
40+
mean_and_max);
41+
}
42+
}
43+
44+
static auto registry =
45+
torch::RegisterOperators().op("torch_cmspepr::accumulate_knn", &accumulate_knn);

‎csrc/accumulate_knn_grad.cpp

+45
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
#include <Python.h>
2+
#include <torch/script.h>
3+
4+
#include "cpu/accumulate_knn_grad_cpu.h"
5+
6+
#ifdef WITH_CUDA
7+
#include "cuda/accumulate_knn_grad_cuda.h"
8+
#endif
9+
10+
11+
#ifdef _WIN32
12+
#ifdef WITH_CUDA
13+
PyMODINIT_FUNC PyInit__accumulate_knn_grad_cuda(void) { return NULL; }
14+
#else
15+
PyMODINIT_FUNC PyInit__accumulate_knn_grad_cpu(void) { return NULL; }
16+
#endif
17+
#endif
18+
19+
std::tuple<torch::Tensor, torch::Tensor>
20+
accumulate_knn_grad(torch::Tensor grad_from_out_features,
21+
torch::Tensor distances,
22+
torch::Tensor features,
23+
torch::Tensor neigh_indices,
24+
torch::Tensor max_feat_indices) {
25+
if (distances.device().is_cuda()) {
26+
#ifdef WITH_CUDA
27+
return accumulate_knn_grad_cuda(grad_from_out_features,
28+
distances,
29+
features,
30+
neigh_indices,
31+
max_feat_indices);
32+
#else
33+
AT_ERROR("Not compiled with CUDA support");
34+
#endif
35+
} else {
36+
return accumulate_knn_grad_cpu(grad_from_out_features,
37+
distances,
38+
features,
39+
neigh_indices,
40+
max_feat_indices);
41+
}
42+
}
43+
44+
static auto registry =
45+
torch::RegisterOperators().op("torch_cmspepr::accumulate_knn_grad", &accumulate_knn_grad);

‎csrc/cmspepr.h

+33
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
#pragma once
2+
3+
#include <torch/extension.h>
4+
5+
std::tuple<torch::Tensor, torch::Tensor>
6+
accumulate_knn(torch::Tensor distances,
7+
torch::Tensor features,
8+
torch::Tensor indices,
9+
int64_t n_moments = 1,
10+
bool mean_and_max = true);
11+
12+
std::tuple<torch::Tensor, torch::Tensor>
13+
accumulate_knn_grad(torch::Tensor grad_from_out_features,
14+
torch::Tensor distances,
15+
torch::Tensor features,
16+
torch::Tensor neigh_indices,
17+
torch::Tensor max_feat_indices);
18+
19+
std::tuple<torch::Tensor, torch::Tensor>
20+
select_knn(torch::Tensor coords,
21+
torch::Tensor row_splits,
22+
torch::Tensor mask,
23+
int64_t n_neighbours,
24+
double max_radius,
25+
int64_t mask_mode = 1);
26+
27+
torch::Tensor
28+
select_knn_grad(torch::Tensor gradDistances,
29+
torch::Tensor indices,
30+
torch::Tensor distances,
31+
torch::Tensor coordinates);
32+
33+

‎csrc/cpu/accumulate_knn_cpu.cpp

+110
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
#include "accumulate_knn_cpu.h"
2+
#include "helpers.h"
3+
#include "utils.h"
4+
5+
#include <torch/extension.h>
6+
#include <string> //size_t, just for helper function
7+
#include <cmath>
8+
//#include <iostream> //remove later DEBUG FIXME
9+
10+
static inline float distanceWeight(const float& distsq) {
11+
return exp(-1. * ACCUMULATE_KNN_EXPONENT * distsq);
12+
}
13+
14+
void compute(const float_t *d_distances,
15+
const float_t *d_feat,
16+
const int32_t *d_idxs,
17+
18+
float_t *d_out_feat,
19+
int32_t *d_out_maxidxs,
20+
21+
size_t n_vert,
22+
size_t n_neigh,
23+
size_t n_feat,
24+
25+
size_t n_out_feat,
26+
27+
size_t n_moments,
28+
bool mean_and_max)
29+
{
30+
for (size_t i_v = 0; i_v < n_vert; i_v++) {
31+
32+
for (size_t i_f = 0; i_f < n_feat; i_f++) {
33+
34+
float t_mean = 0;
35+
float t_max = 0;
36+
int max_i_n_gidx = 0;
37+
38+
for (size_t i_n = 0; i_n < n_neigh; i_n++) {
39+
40+
int nidx = d_idxs[I2D(i_v, i_n, n_neigh)];
41+
if (nidx < 0) continue;
42+
43+
float vnf = d_feat[I2D(nidx, i_f, n_feat)];
44+
float distsq = d_distances[I2D(i_v, i_n, n_neigh)];
45+
float wfeat = vnf * distanceWeight(distsq);
46+
47+
t_mean += wfeat;
48+
49+
if (mean_and_max && (wfeat >= t_max || !i_n)) {
50+
max_i_n_gidx = nidx;
51+
t_max = wfeat;
52+
}
53+
}
54+
55+
t_mean /= (float)n_neigh;
56+
d_out_feat[I2D(i_v, i_f, n_out_feat)] = t_mean;
57+
if (mean_and_max) {
58+
d_out_maxidxs[I2D(i_v, i_f, n_feat)] = max_i_n_gidx;
59+
d_out_feat[I2D(i_v, i_f + n_feat, n_out_feat)] = t_max;
60+
}
61+
62+
}
63+
}
64+
}
65+
66+
std::tuple<torch::Tensor, torch::Tensor>
67+
accumulate_knn_cpu(torch::Tensor distances,
68+
torch::Tensor features,
69+
torch::Tensor indices,
70+
int n_moments,
71+
bool mean_and_max)
72+
{
73+
const auto n_vert = distances.size(0);
74+
const auto n_neigh = indices.size(1);
75+
const auto n_coords = distances.size(1);
76+
const auto n_feat = features.size(1);
77+
78+
assert(n_vert == indices.size(0) && n_vert == features.size(0));
79+
assert(n_neigh == distances.size(1));
80+
81+
int64_t n_out_feat = n_feat;
82+
if (mean_and_max) {
83+
n_out_feat *= 2; }
84+
85+
auto output_feat_tensor = torch::zeros({ n_vert,n_out_feat },
86+
torch::TensorOptions().dtype(torch::kFloat32));
87+
auto output_max_idxs_tensor = torch::zeros({ n_vert,n_feat },
88+
torch::TensorOptions().dtype(torch::kInt32));
89+
90+
auto distances_data = distances.data_ptr<float_t>();
91+
auto features_data = features.data_ptr<float_t>();
92+
auto indices_data = indices.data_ptr<int32_t>();
93+
94+
auto output_feat_tensor_data = output_feat_tensor.data_ptr<float_t>();
95+
auto output_max_idxs_data = output_max_idxs_tensor.data_ptr<int32_t>();
96+
97+
compute(distances_data,
98+
features_data,
99+
indices_data,
100+
output_feat_tensor_data,
101+
output_max_idxs_data,
102+
n_vert,
103+
n_neigh,
104+
n_feat,
105+
n_out_feat,
106+
n_moments,
107+
mean_and_max);
108+
109+
return std::make_tuple(output_feat_tensor, output_max_idxs_tensor);
110+
}

‎csrc/cpu/accumulate_knn_cpu.h

+12
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
#pragma once
2+
3+
#define ACCUMULATE_KNN_EXPONENT 1.
4+
5+
#include <torch/extension.h>
6+
7+
std::tuple<torch::Tensor, torch::Tensor>
8+
accumulate_knn_cpu(torch::Tensor distances,
9+
torch::Tensor features,
10+
torch::Tensor indices,
11+
int n_moments = 1,
12+
bool mean_and_max = true);

0 commit comments

Comments
 (0)