diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 66b13b6710..b48eff36db 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -119,28 +119,66 @@ rapids_cmake_install_lib_dir( lib_dir ) include(CPack) +file(GLOB_RECURSE RAFT_DISTANCE_SOURCES "src/distance/specializations/*.cu") +file(GLOB_RECURSE RAFT_NN_SOURCES "src/nn/specializations/*.cu" ) + +add_library(raft_distance SHARED ${RAFT_DISTANCE_SOURCES}) +add_library(raft::raft_distance ALIAS raft_distance) + +add_library(raft_nn SHARED ${RAFT_NN_SOURCES}) +add_library(raft::raft_nn ALIAS raft_nn) + add_library(raft INTERFACE) add_library(raft::raft ALIAS raft) target_include_directories(raft INTERFACE "$" "$") -target_link_libraries(raft -INTERFACE - CUDA::cublas - CUDA::curand - CUDA::cusolver - CUDA::cudart - CUDA::cusparse - rmm::rmm - cuco::cuco - ) +target_include_directories(raft_distance PUBLIC "$" + "$") + +target_include_directories(raft_nn PUBLIC "$" + "$") + +set(RAFT_LINK_LIBRARIES + CUDA::cublas + CUDA::curand + CUDA::cusolver + CUDA::cudart + CUDA::cusparse + rmm::rmm + cuco::cuco + ) + +target_link_libraries(raft INTERFACE ${RAFT_LINK_LIBRARIES}) +target_link_libraries(raft_distance PUBLIC ${RAFT_LINK_LIBRARIES}) +target_link_libraries(raft_nn PUBLIC ${RAFT_LINK_LIBRARIES} FAISS::FAISS) + +target_compile_options(raft_distance + PRIVATE "$<$:${RAFT_CXX_FLAGS}>" + "$<$:${RAFT_CUDA_FLAGS}>" + ) + + +target_compile_options(raft_nn + PRIVATE "$<$:${RAFT_CXX_FLAGS}>" + "$<$:${RAFT_CUDA_FLAGS}>" + ) +target_compile_features(raft_distance PUBLIC cxx_std_17 $) +target_compile_features(raft_nn PUBLIC cxx_std_17 $) target_compile_features(raft INTERFACE cxx_std_17 $) +install(TARGETS raft_distance + DESTINATION ${lib_dir} + EXPORT raft-exports) + +install(TARGETS raft_nn + DESTINATION ${lib_dir} + EXPORT raft-exports) + install(TARGETS raft DESTINATION ${lib_dir} - EXPORT raft-exports - ) + EXPORT raft-exports) include(GNUInstallDirs) install(DIRECTORY include/raft/ @@ -149,23 +187,22 @@ install(DIRECTORY include/raft/ # Temporary install of raft.hpp while the file is removed install(FILES include/raft.hpp - DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/raft - ) + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/raft) ############################################################################## # - install export ----------------------------------------------------------- set(doc_string [=[ -Provide targets for the RAFT: RAPIDS Analytics Frameworks Toolset. - -RAFT is a repository containining shared utilities, mathematical operations -and common functions for the analytics components of RAPIDS. +Provide targets for the RAFT: RAPIDS Analytics Framework Toolkit. +RAPIDS Analytics Framework Toolkit contains shared representations, +mathematical computational primitives, and utilities that accelerate +building analytics and data science algorithms in the RAPIDS ecosystem. ]=]) rapids_export(INSTALL raft EXPORT_SET raft-exports - GLOBAL_TARGETS raft # since we can't hook into EXPORT SETS + GLOBAL_TARGETS raft raft_distance# since we can't hook into EXPORT SETS NAMESPACE raft:: DOCUMENTATION doc_string ) @@ -175,7 +212,7 @@ and common functions for the analytics components of RAPIDS. rapids_export(BUILD raft EXPORT_SET raft-exports - GLOBAL_TARGETS raft # since we can't hook into EXPORT SETS + GLOBAL_TARGETS raft raft_distance raft_nn# since we can't hook into EXPORT SETS LANGUAGES CUDA DOCUMENTATION doc_string NAMESPACE raft:: diff --git a/cpp/include/raft/distance/specializations.hpp b/cpp/include/raft/distance/specializations.hpp new file mode 100644 index 0000000000..e70943e731 --- /dev/null +++ b/cpp/include/raft/distance/specializations.hpp @@ -0,0 +1,19 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include \ No newline at end of file diff --git a/cpp/include/raft/distance/specializations/detail/canberra.hpp b/cpp/include/raft/distance/specializations/detail/canberra.hpp new file mode 100644 index 0000000000..2e71685532 --- /dev/null +++ b/cpp/include/raft/distance/specializations/detail/canberra.hpp @@ -0,0 +1,67 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +namespace raft { +namespace distance { +namespace detail { +extern template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +extern template void distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +extern template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/include/raft/distance/specializations/detail/chebyshev.hpp b/cpp/include/raft/distance/specializations/detail/chebyshev.hpp new file mode 100644 index 0000000000..dc03e047be --- /dev/null +++ b/cpp/include/raft/distance/specializations/detail/chebyshev.hpp @@ -0,0 +1,66 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace raft { +namespace distance { +namespace detail { +extern template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +extern template void distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +extern template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/include/raft/distance/specializations/detail/correlation.hpp b/cpp/include/raft/distance/specializations/detail/correlation.hpp new file mode 100644 index 0000000000..2e7683ab10 --- /dev/null +++ b/cpp/include/raft/distance/specializations/detail/correlation.hpp @@ -0,0 +1,68 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace raft { +namespace distance { +namespace detail { +extern template void +distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +extern template void +distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +extern template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/include/raft/distance/specializations/detail/cosine.hpp b/cpp/include/raft/distance/specializations/detail/cosine.hpp new file mode 100644 index 0000000000..b47d294645 --- /dev/null +++ b/cpp/include/raft/distance/specializations/detail/cosine.hpp @@ -0,0 +1,68 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace raft { +namespace distance { +namespace detail { +extern template void +distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +extern template void +distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +extern template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/include/raft/distance/specializations/detail/hamming_unexpanded.hpp b/cpp/include/raft/distance/specializations/detail/hamming_unexpanded.hpp new file mode 100644 index 0000000000..29a4ca03d9 --- /dev/null +++ b/cpp/include/raft/distance/specializations/detail/hamming_unexpanded.hpp @@ -0,0 +1,68 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace raft { +namespace distance { +namespace detail { +extern template void +distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +extern template void +distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +extern template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/include/raft/distance/specializations/detail/hellinger_expanded.hpp b/cpp/include/raft/distance/specializations/detail/hellinger_expanded.hpp new file mode 100644 index 0000000000..264003ec0e --- /dev/null +++ b/cpp/include/raft/distance/specializations/detail/hellinger_expanded.hpp @@ -0,0 +1,68 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace raft { +namespace distance { +namespace detail { +extern template void +distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +extern template void +distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +extern template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/include/raft/distance/specializations/detail/jensen_shannon.hpp b/cpp/include/raft/distance/specializations/detail/jensen_shannon.hpp new file mode 100644 index 0000000000..3135a4c579 --- /dev/null +++ b/cpp/include/raft/distance/specializations/detail/jensen_shannon.hpp @@ -0,0 +1,66 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace raft { +namespace distance { +namespace detail { +extern template void +distance(const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +extern template void +distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); +extern template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/include/raft/distance/specializations/detail/kl_divergence.hpp b/cpp/include/raft/distance/specializations/detail/kl_divergence.hpp new file mode 100644 index 0000000000..207fca6bc2 --- /dev/null +++ b/cpp/include/raft/distance/specializations/detail/kl_divergence.hpp @@ -0,0 +1,67 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace raft { +namespace distance { +namespace detail { +extern template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +extern template void +distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +extern template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/include/raft/distance/specializations/detail/l1.hpp b/cpp/include/raft/distance/specializations/detail/l1.hpp new file mode 100644 index 0000000000..e8eddfe1e4 --- /dev/null +++ b/cpp/include/raft/distance/specializations/detail/l1.hpp @@ -0,0 +1,65 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace raft { +namespace distance { +namespace detail { +extern template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +extern template void distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +extern template void distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/include/raft/distance/specializations/detail/l2_expanded.hpp b/cpp/include/raft/distance/specializations/detail/l2_expanded.hpp new file mode 100644 index 0000000000..db37b8db9f --- /dev/null +++ b/cpp/include/raft/distance/specializations/detail/l2_expanded.hpp @@ -0,0 +1,66 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace raft { +namespace distance { +namespace detail { +extern template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +extern template void +distance(const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +extern template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/include/raft/distance/specializations/detail/l2_sqrt_expanded.hpp b/cpp/include/raft/distance/specializations/detail/l2_sqrt_expanded.hpp new file mode 100644 index 0000000000..ac23c9c357 --- /dev/null +++ b/cpp/include/raft/distance/specializations/detail/l2_sqrt_expanded.hpp @@ -0,0 +1,68 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace raft { +namespace distance { +namespace detail { +extern template void +distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +extern template void +distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +extern template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/include/raft/distance/specializations/detail/l2_sqrt_unexpanded.hpp b/cpp/include/raft/distance/specializations/detail/l2_sqrt_unexpanded.hpp new file mode 100644 index 0000000000..1e38575fbf --- /dev/null +++ b/cpp/include/raft/distance/specializations/detail/l2_sqrt_unexpanded.hpp @@ -0,0 +1,68 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace raft { +namespace distance { +namespace detail { +extern template void +distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +extern template void +distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +extern template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/include/raft/distance/specializations/detail/l2_unexpanded.hpp b/cpp/include/raft/distance/specializations/detail/l2_unexpanded.hpp new file mode 100644 index 0000000000..035c9ef693 --- /dev/null +++ b/cpp/include/raft/distance/specializations/detail/l2_unexpanded.hpp @@ -0,0 +1,67 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace raft { +namespace distance { +namespace detail { +extern template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +extern template void +distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +extern template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/include/raft/distance/specializations/detail/lp_unexpanded.hpp b/cpp/include/raft/distance/specializations/detail/lp_unexpanded.hpp new file mode 100644 index 0000000000..83eda5f07b --- /dev/null +++ b/cpp/include/raft/distance/specializations/detail/lp_unexpanded.hpp @@ -0,0 +1,67 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace raft { +namespace distance { +namespace detail { +extern template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +extern template void +distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +extern template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/include/raft/distance/specializations/distance.hpp b/cpp/include/raft/distance/specializations/distance.hpp new file mode 100644 index 0000000000..a57d6f49a5 --- /dev/null +++ b/cpp/include/raft/distance/specializations/distance.hpp @@ -0,0 +1,32 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include diff --git a/cpp/include/raft/sparse/selection/detail/knn_graph.cuh b/cpp/include/raft/sparse/selection/detail/knn_graph.cuh index 83cb23f513..3fad21f307 100644 --- a/cpp/include/raft/sparse/selection/detail/knn_graph.cuh +++ b/cpp/include/raft/sparse/selection/detail/knn_graph.cuh @@ -104,7 +104,7 @@ void knn_graph(const handle_t& handle, raft::sparse::COO& out, int c = 15) { - int k = build_k(m, c); + size_t k = build_k(m, c); auto stream = handle.get_stream(); @@ -120,7 +120,7 @@ void knn_graph(const handle_t& handle, std::vector inputs; inputs.push_back(const_cast(X)); - std::vector sizes; + std::vector sizes; sizes.push_back(m); // This is temporary. Once faiss is updated, we should be able to @@ -128,19 +128,19 @@ void knn_graph(const handle_t& handle, rmm::device_uvector int64_indices(nnz, stream); uint32_t knn_start = curTimeMillis(); - raft::spatial::knn::brute_force_knn(handle, - inputs, - sizes, - n, - const_cast(X), - m, - int64_indices.data(), - data.data(), - k, - true, - true, - nullptr, - metric); + raft::spatial::knn::brute_force_knn(handle, + inputs, + sizes, + n, + const_cast(X), + m, + int64_indices.data(), + data.data(), + k, + true, + true, + nullptr, + metric); // convert from current knn's 64-bit to 32-bit. conv_indices(int64_indices.data(), indices.data(), nnz, stream); @@ -152,4 +152,4 @@ void knn_graph(const handle_t& handle, }; // namespace detail }; // namespace selection }; // namespace sparse -}; // end namespace raft \ No newline at end of file +}; // end namespace raft diff --git a/cpp/include/raft/sparse/selection/knn_graph.hpp b/cpp/include/raft/sparse/selection/knn_graph.hpp index 7af452541f..96ce02e06a 100644 --- a/cpp/include/raft/sparse/selection/knn_graph.hpp +++ b/cpp/include/raft/sparse/selection/knn_graph.hpp @@ -17,6 +17,7 @@ #pragma once #include +#include #include #include @@ -44,8 +45,8 @@ namespace selection { template void knn_graph(const handle_t& handle, const value_t* X, - size_t m, - size_t n, + std::size_t m, + std::size_t n, raft::distance::DistanceType metric, raft::sparse::COO& out, int c = 15) @@ -55,4 +56,4 @@ void knn_graph(const handle_t& handle, }; // namespace selection }; // namespace sparse -}; // end namespace raft \ No newline at end of file +}; // end namespace raft diff --git a/cpp/include/raft/spatial/knn/ball_cover.hpp b/cpp/include/raft/spatial/knn/ball_cover.hpp index cb2b9e99cd..39f5845794 100644 --- a/cpp/include/raft/spatial/knn/ball_cover.hpp +++ b/cpp/include/raft/spatial/knn/ball_cover.hpp @@ -34,10 +34,10 @@ void rbc_build_index(const raft::handle_t& handle, { ASSERT(index.n == 2, "Random ball cover currently only works in 2-dimensions"); if (index.metric == raft::distance::DistanceType::Haversine) { - detail::rbc_build_index(handle, index, detail::HaversineFunc()); + detail::rbc_build_index(handle, index, detail::HaversineFunc()); } else if (index.metric == raft::distance::DistanceType::L2SqrtExpanded || index.metric == raft::distance::DistanceType::L2SqrtUnexpanded) { - detail::rbc_build_index(handle, index, detail::EuclideanFunc()); + detail::rbc_build_index(handle, index, detail::EuclideanFunc()); } else { RAFT_FAIL("Metric not support"); } @@ -84,12 +84,24 @@ void rbc_all_knn_query(const raft::handle_t& handle, { ASSERT(index.n == 2, "Random ball cover currently only works in 2-dimensions"); if (index.metric == raft::distance::DistanceType::Haversine) { - detail::rbc_all_knn_query( - handle, index, k, inds, dists, detail::HaversineFunc(), perform_post_filtering, weight); + detail::rbc_all_knn_query(handle, + index, + k, + inds, + dists, + detail::HaversineFunc(), + perform_post_filtering, + weight); } else if (index.metric == raft::distance::DistanceType::L2SqrtExpanded || index.metric == raft::distance::DistanceType::L2SqrtUnexpanded) { - detail::rbc_all_knn_query( - handle, index, k, inds, dists, detail::EuclideanFunc(), perform_post_filtering, weight); + detail::rbc_all_knn_query(handle, + index, + k, + inds, + dists, + detail::EuclideanFunc(), + perform_post_filtering, + weight); } else { RAFT_FAIL("Metric not supported"); } @@ -146,7 +158,7 @@ void rbc_knn_query(const raft::handle_t& handle, n_query_pts, inds, dists, - detail::HaversineFunc(), + detail::HaversineFunc(), perform_post_filtering, weight); } else if (index.metric == raft::distance::DistanceType::L2SqrtExpanded || @@ -158,7 +170,7 @@ void rbc_knn_query(const raft::handle_t& handle, n_query_pts, inds, dists, - detail::EuclideanFunc(), + detail::EuclideanFunc(), perform_post_filtering, weight); } else { diff --git a/cpp/include/raft/spatial/knn/detail/ball_cover/common.cuh b/cpp/include/raft/spatial/knn/detail/ball_cover/common.cuh index 181dad1a90..4b3065a0f3 100644 --- a/cpp/include/raft/spatial/knn/detail/ball_cover/common.cuh +++ b/cpp/include/raft/spatial/knn/detail/ball_cover/common.cuh @@ -38,21 +38,31 @@ struct NNComp { } }; -struct HaversineFunc { - template +template +struct DistFunc { + virtual __device__ __host__ __forceinline__ value_t operator()(const value_t* a, + const value_t* b, + const value_int n_dims) + { + return -1; + }; +}; + +template +struct HaversineFunc : public DistFunc { __device__ __host__ __forceinline__ value_t operator()(const value_t* a, const value_t* b, - const value_int n_dims) + const value_int n_dims) override { return raft::spatial::knn::detail::compute_haversine(a[0], b[0], a[1], b[1]); } }; -struct EuclideanFunc { - template +template +struct EuclideanFunc : public DistFunc { __device__ __host__ __forceinline__ value_t operator()(const value_t* a, const value_t* b, - const value_int n_dims) + const value_int n_dims) override { value_t sum_sq = 0; for (value_int i = 0; i < n_dims; ++i) { diff --git a/cpp/include/raft/spatial/knn/detail/ball_cover/registers.cuh b/cpp/include/raft/spatial/knn/detail/ball_cover/registers.cuh index 5d28258f7a..32d8068834 100644 --- a/cpp/include/raft/spatial/knn/detail/ball_cover/registers.cuh +++ b/cpp/include/raft/spatial/knn/detail/ball_cover/registers.cuh @@ -470,7 +470,7 @@ void rbc_low_dim_pass_one(const raft::handle_t& handle, value_int k, const value_idx* R_knn_inds, const value_t* R_knn_dists, - dist_func dfunc, + dist_func& dfunc, value_idx* inds, value_t* dists, float weight, @@ -601,7 +601,7 @@ void rbc_low_dim_pass_two(const raft::handle_t& handle, value_int k, const value_idx* R_knn_inds, const value_t* R_knn_dists, - dist_func dfunc, + dist_func& dfunc, value_idx* inds, value_t* dists, float weight, @@ -611,7 +611,7 @@ void rbc_low_dim_pass_two(const raft::handle_t& handle, rmm::device_uvector bitset(bitset_size * index.m, handle.get_stream()); - perform_post_filter_registers + perform_post_filter_registers <<>>( index.get_X(), index.n, diff --git a/cpp/include/raft/spatial/knn/detail/knn_brute_force_faiss.cuh b/cpp/include/raft/spatial/knn/detail/knn_brute_force_faiss.cuh index 12b7124773..68590d3d7d 100644 --- a/cpp/include/raft/spatial/knn/detail/knn_brute_force_faiss.cuh +++ b/cpp/include/raft/spatial/knn/detail/knn_brute_force_faiss.cuh @@ -216,16 +216,16 @@ inline void knn_merge_parts(value_t* inK, * @param[in] metric corresponds to the raft::distance::DistanceType enum (default is L2Expanded) * @param[in] metricArg metric argument to use. Corresponds to the p arg for lp norm */ -template +template void brute_force_knn_impl( const raft::handle_t& handle, - std::vector& input, + std::vector& input, std::vector& sizes, IntType D, - float* search_items, + value_t* search_items, IntType n, IdxType* res_I, - float* res_D, + value_t* res_D, IntType k, bool rowMajorIndex = true, bool rowMajorQuery = true, @@ -254,14 +254,14 @@ void brute_force_knn_impl( } // perform preprocessing - std::unique_ptr> query_metric_processor = - create_processor(metric, n, D, k, rowMajorQuery, userStream); + std::unique_ptr> query_metric_processor = + create_processor(metric, n, D, k, rowMajorQuery, userStream); query_metric_processor->preprocess(search_items); - std::vector>> metric_processors(input.size()); + std::vector>> metric_processors(input.size()); for (size_t i = 0; i < input.size(); i++) { metric_processors[i] = - create_processor(metric, sizes[i], D, k, rowMajorQuery, userStream); + create_processor(metric, sizes[i], D, k, rowMajorQuery, userStream); metric_processors[i]->preprocess(input[i]); } @@ -271,10 +271,10 @@ void brute_force_knn_impl( rmm::device_uvector trans(id_ranges->size(), userStream); raft::update_device(trans.data(), id_ranges->data(), id_ranges->size(), userStream); - rmm::device_uvector all_D(0, userStream); + rmm::device_uvector all_D(0, userStream); rmm::device_uvector all_I(0, userStream); - float* out_D = res_D; + value_t* out_D = res_D; IdxType* out_I = res_I; if (input.size() > 1) { @@ -289,7 +289,7 @@ void brute_force_knn_impl( handle.wait_stream_pool_on_stream(); for (size_t i = 0; i < input.size(); i++) { - float* out_d_ptr = out_D + (i * k * n); + value_t* out_d_ptr = out_D + (i * k * n); IdxType* out_i_ptr = out_I + (i * k * n); auto stream = handle.get_next_usable_stream(i); diff --git a/cpp/include/raft/spatial/knn/knn.hpp b/cpp/include/raft/spatial/knn/knn.hpp index e1e1eac248..b29c4cc51c 100644 --- a/cpp/include/raft/spatial/knn/knn.hpp +++ b/cpp/include/raft/spatial/knn/knn.hpp @@ -123,20 +123,21 @@ inline void select_k(value_t* inK, * @param[in] translations starting offsets for partitions. should be the same size * as input vector. */ -inline void brute_force_knn(raft::handle_t const& handle, - std::vector& input, - std::vector& sizes, - int D, - float* search_items, - int n, - int64_t* res_I, - float* res_D, - int k, - bool rowMajorIndex = true, - bool rowMajorQuery = true, - std::vector* translations = nullptr, - distance::DistanceType metric = distance::DistanceType::L2Expanded, - float metric_arg = 2.0f) +template +void brute_force_knn(raft::handle_t const& handle, + std::vector& input, + std::vector& sizes, + value_int D, + value_t* search_items, + value_int n, + value_idx* res_I, + value_t* res_D, + value_int k, + bool rowMajorIndex = true, + bool rowMajorQuery = true, + std::vector* translations = nullptr, + distance::DistanceType metric = distance::DistanceType::L2Unexpanded, + float metric_arg = 2.0f) { ASSERT(input.size() == sizes.size(), "input and sizes vectors must be the same size"); diff --git a/cpp/include/raft/spatial/knn/specializations.hpp b/cpp/include/raft/spatial/knn/specializations.hpp new file mode 100644 index 0000000000..663e77c6a0 --- /dev/null +++ b/cpp/include/raft/spatial/knn/specializations.hpp @@ -0,0 +1,21 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include diff --git a/cpp/include/raft/spatial/knn/specializations/ball_cover.hpp b/cpp/include/raft/spatial/knn/specializations/ball_cover.hpp new file mode 100644 index 0000000000..dad0b3f1ee --- /dev/null +++ b/cpp/include/raft/spatial/knn/specializations/ball_cover.hpp @@ -0,0 +1,55 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include +#include +#include + +namespace raft { +namespace spatial { +namespace knn { +extern template class BallCoverIndex; +extern template class BallCoverIndex; + +extern template void rbc_build_index( + const raft::handle_t& handle, BallCoverIndex& index); + +extern template void rbc_knn_query( + const raft::handle_t& handle, + BallCoverIndex& index, + std::uint32_t k, + const float* query, + std::uint32_t n_query_pts, + std::int64_t* inds, + float* dists, + bool perform_post_filtering, + float weight); + +extern template void rbc_all_knn_query( + const raft::handle_t& handle, + BallCoverIndex& index, + std::uint32_t k, + std::int64_t* inds, + float* dists, + bool perform_post_filtering, + float weight); +}; // namespace knn +}; // namespace spatial +}; // namespace raft diff --git a/cpp/include/raft/spatial/knn/specializations/detail/ball_cover_lowdim.hpp b/cpp/include/raft/spatial/knn/specializations/detail/ball_cover_lowdim.hpp new file mode 100644 index 0000000000..d0e4813332 --- /dev/null +++ b/cpp/include/raft/spatial/knn/specializations/detail/ball_cover_lowdim.hpp @@ -0,0 +1,56 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include + +namespace raft { +namespace spatial { +namespace knn { +namespace detail { + +extern template void rbc_low_dim_pass_one( + const raft::handle_t& handle, + BallCoverIndex& index, + const float* query, + const std::uint32_t n_query_rows, + std::uint32_t k, + const std::int64_t* R_knn_inds, + const float* R_knn_dists, + DistFunc& dfunc, + std::int64_t* inds, + float* dists, + float weight, + std::uint32_t* dists_counter); + +extern template void rbc_low_dim_pass_two( + const raft::handle_t& handle, + BallCoverIndex& index, + const float* query, + const std::uint32_t n_query_rows, + std::uint32_t k, + const std::int64_t* R_knn_inds, + const float* R_knn_dists, + DistFunc& dfunc, + std::int64_t* inds, + float* dists, + float weight, + std::uint32_t* post_dists_counter); +}; // namespace detail +}; // namespace knn +}; // namespace spatial +}; // namespace raft \ No newline at end of file diff --git a/cpp/include/raft/spatial/knn/specializations/fused_l2_knn.hpp b/cpp/include/raft/spatial/knn/specializations/fused_l2_knn.hpp new file mode 100644 index 0000000000..961351d734 --- /dev/null +++ b/cpp/include/raft/spatial/knn/specializations/fused_l2_knn.hpp @@ -0,0 +1,80 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +namespace raft { +namespace spatial { +namespace knn { +namespace detail { + +extern template void fusedL2Knn(size_t D, + long* out_inds, + float* out_dists, + const float* index, + const float* query, + size_t n_index_rows, + size_t n_query_rows, + int k, + bool rowMajorIndex, + bool rowMajorQuery, + cudaStream_t stream, + raft::distance::DistanceType metric); + +extern template void fusedL2Knn(size_t D, + long* out_inds, + float* out_dists, + const float* index, + const float* query, + size_t n_index_rows, + size_t n_query_rows, + int k, + bool rowMajorIndex, + bool rowMajorQuery, + cudaStream_t stream, + raft::distance::DistanceType metric); + +extern template void fusedL2Knn(size_t D, + int* out_inds, + float* out_dists, + const float* index, + const float* query, + size_t n_index_rows, + size_t n_query_rows, + int k, + bool rowMajorIndex, + bool rowMajorQuery, + cudaStream_t stream, + raft::distance::DistanceType metric); + +extern template void fusedL2Knn(size_t D, + int* out_inds, + float* out_dists, + const float* index, + const float* query, + size_t n_index_rows, + size_t n_query_rows, + int k, + bool rowMajorIndex, + bool rowMajorQuery, + cudaStream_t stream, + raft::distance::DistanceType metric); + +}; // namespace detail +}; // namespace knn +}; // namespace spatial +}; // namespace raft diff --git a/cpp/include/raft/spatial/knn/specializations/knn.hpp b/cpp/include/raft/spatial/knn/specializations/knn.hpp new file mode 100644 index 0000000000..bd8673af39 --- /dev/null +++ b/cpp/include/raft/spatial/knn/specializations/knn.hpp @@ -0,0 +1,55 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace raft { +namespace spatial { +namespace knn { +extern template void brute_force_knn(raft::handle_t const& handle, + std::vector& input, + std::vector& sizes, + int D, + float* search_items, + int n, + long* res_I, + float* res_D, + int k, + bool rowMajorIndex, + bool rowMajorQuery, + std::vector* translations, + distance::DistanceType metric, + float metric_arg); + +extern template void brute_force_knn(raft::handle_t const& handle, + std::vector& input, + std::vector& sizes, + unsigned int D, + float* search_items, + unsigned int n, + long* res_I, + float* res_D, + unsigned int k, + bool rowMajorIndex, + bool rowMajorQuery, + std::vector* translations, + distance::DistanceType metric, + float metric_arg); +}; // namespace knn +}; // namespace spatial +}; // namespace raft diff --git a/cpp/src/distance/specializations/detail/canberra.cu b/cpp/src/distance/specializations/detail/canberra.cu new file mode 100644 index 0000000000..b2dd993ab7 --- /dev/null +++ b/cpp/src/distance/specializations/detail/canberra.cu @@ -0,0 +1,63 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft { +namespace distance { +namespace detail { +template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +template void distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +template void distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/src/distance/specializations/detail/chebyshev.cu b/cpp/src/distance/specializations/detail/chebyshev.cu new file mode 100644 index 0000000000..ab310515bd --- /dev/null +++ b/cpp/src/distance/specializations/detail/chebyshev.cu @@ -0,0 +1,63 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft { +namespace distance { +namespace detail { +template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +template void distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +template void distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/src/distance/specializations/detail/correlation.cu b/cpp/src/distance/specializations/detail/correlation.cu new file mode 100644 index 0000000000..04b9e5bf69 --- /dev/null +++ b/cpp/src/distance/specializations/detail/correlation.cu @@ -0,0 +1,65 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft { +namespace distance { +namespace detail { +template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +template void +distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/src/distance/specializations/detail/cosine.cu b/cpp/src/distance/specializations/detail/cosine.cu new file mode 100644 index 0000000000..bc19599511 --- /dev/null +++ b/cpp/src/distance/specializations/detail/cosine.cu @@ -0,0 +1,64 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft { +namespace distance { +namespace detail { +template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +template void distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/src/distance/specializations/detail/hamming_unexpanded.cu b/cpp/src/distance/specializations/detail/hamming_unexpanded.cu new file mode 100644 index 0000000000..e5e66e85bd --- /dev/null +++ b/cpp/src/distance/specializations/detail/hamming_unexpanded.cu @@ -0,0 +1,65 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft { +namespace distance { +namespace detail { +template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +template void +distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/src/distance/specializations/detail/hellinger_expanded.cu b/cpp/src/distance/specializations/detail/hellinger_expanded.cu new file mode 100644 index 0000000000..fa9b8e14d6 --- /dev/null +++ b/cpp/src/distance/specializations/detail/hellinger_expanded.cu @@ -0,0 +1,64 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft { +namespace distance { +namespace detail { +template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +template void +distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); +template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/src/distance/specializations/detail/jensen_shannon.cu b/cpp/src/distance/specializations/detail/jensen_shannon.cu new file mode 100644 index 0000000000..37f1f81fb1 --- /dev/null +++ b/cpp/src/distance/specializations/detail/jensen_shannon.cu @@ -0,0 +1,64 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft { +namespace distance { +namespace detail { +template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +template void distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/src/distance/specializations/detail/kl_divergence.cu b/cpp/src/distance/specializations/detail/kl_divergence.cu new file mode 100644 index 0000000000..f6412cdefd --- /dev/null +++ b/cpp/src/distance/specializations/detail/kl_divergence.cu @@ -0,0 +1,64 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft { +namespace distance { +namespace detail { +template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +template void distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/src/distance/specializations/detail/l1.cu b/cpp/src/distance/specializations/detail/l1.cu new file mode 100644 index 0000000000..5df9c9ece6 --- /dev/null +++ b/cpp/src/distance/specializations/detail/l1.cu @@ -0,0 +1,63 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft { +namespace distance { +namespace detail { +template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +template void distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +template void distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/src/distance/specializations/detail/l2_expanded.cu b/cpp/src/distance/specializations/detail/l2_expanded.cu new file mode 100644 index 0000000000..1b122ca331 --- /dev/null +++ b/cpp/src/distance/specializations/detail/l2_expanded.cu @@ -0,0 +1,64 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft { +namespace distance { +namespace detail { +template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +template void distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/src/distance/specializations/detail/l2_sqrt_expanded.cu b/cpp/src/distance/specializations/detail/l2_sqrt_expanded.cu new file mode 100644 index 0000000000..f87d08b94b --- /dev/null +++ b/cpp/src/distance/specializations/detail/l2_sqrt_expanded.cu @@ -0,0 +1,64 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft { +namespace distance { +namespace detail { +template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +template void distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/src/distance/specializations/detail/l2_sqrt_unexpanded.cu b/cpp/src/distance/specializations/detail/l2_sqrt_unexpanded.cu new file mode 100644 index 0000000000..7067cc9015 --- /dev/null +++ b/cpp/src/distance/specializations/detail/l2_sqrt_unexpanded.cu @@ -0,0 +1,64 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft { +namespace distance { +namespace detail { +template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +template void distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/src/distance/specializations/detail/l2_unexpanded.cu b/cpp/src/distance/specializations/detail/l2_unexpanded.cu new file mode 100644 index 0000000000..bdd57b13b2 --- /dev/null +++ b/cpp/src/distance/specializations/detail/l2_unexpanded.cu @@ -0,0 +1,64 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft { +namespace distance { +namespace detail { +template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +template void distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/src/distance/specializations/detail/lp_unexpanded.cu b/cpp/src/distance/specializations/detail/lp_unexpanded.cu new file mode 100644 index 0000000000..a3c15e498d --- /dev/null +++ b/cpp/src/distance/specializations/detail/lp_unexpanded.cu @@ -0,0 +1,64 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft { +namespace distance { +namespace detail { +template void distance( + const float* x, + const float* y, + float* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +template void distance( + const double* x, + const double* y, + double* dist, + int m, + int n, + int k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + double metric_arg); + +template void +distance( + const float* x, + const float* y, + float* dist, + std::uint32_t m, + std::uint32_t n, + std::uint32_t k, + void* workspace, + std::size_t worksize, + cudaStream_t stream, + bool isRowMajor, + float metric_arg); + +} // namespace detail +} // namespace distance +} // namespace raft diff --git a/cpp/src/nn/specializations/ball_cover.cu b/cpp/src/nn/specializations/ball_cover.cu new file mode 100644 index 0000000000..656aec3323 --- /dev/null +++ b/cpp/src/nn/specializations/ball_cover.cu @@ -0,0 +1,58 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include + +// Ignore upstream specializations to avoid unnecessary recompiling +#include +#include +#include +#include + +namespace raft { +namespace spatial { +namespace knn { +template class BallCoverIndex; +template class BallCoverIndex; + +template void rbc_build_index( + const raft::handle_t& handle, BallCoverIndex& index); + +template void rbc_knn_query( + const raft::handle_t& handle, + BallCoverIndex& index, + std::uint32_t k, + const float* query, + std::uint32_t n_query_pts, + std::int64_t* inds, + float* dists, + bool perform_post_filtering, + float weight); + +template void rbc_all_knn_query( + const raft::handle_t& handle, + BallCoverIndex& index, + std::uint32_t k, + std::int64_t* inds, + float* dists, + bool perform_post_filtering, + float weight); + +}; // namespace knn +}; // namespace spatial +}; // namespace raft diff --git a/cpp/src/nn/specializations/detail/ball_cover_lowdim.cu b/cpp/src/nn/specializations/detail/ball_cover_lowdim.cu new file mode 100644 index 0000000000..dea7fe8d41 --- /dev/null +++ b/cpp/src/nn/specializations/detail/ball_cover_lowdim.cu @@ -0,0 +1,55 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +namespace raft { +namespace spatial { +namespace knn { +namespace detail { + +template void rbc_low_dim_pass_one( + const raft::handle_t& handle, + BallCoverIndex& index, + const float* query, + const std::uint32_t n_query_rows, + std::uint32_t k, + const std::int64_t* R_knn_inds, + const float* R_knn_dists, + DistFunc& dfunc, + std::int64_t* inds, + float* dists, + float weight, + std::uint32_t* dists_counter); + +template void rbc_low_dim_pass_two( + const raft::handle_t& handle, + BallCoverIndex& index, + const float* query, + const std::uint32_t n_query_rows, + std::uint32_t k, + const std::int64_t* R_knn_inds, + const float* R_knn_dists, + DistFunc& dfunc, + std::int64_t* inds, + float* dists, + float weight, + std::uint32_t* post_dists_counter); +}; // namespace detail +}; // namespace knn +}; // namespace spatial +}; // namespace raft \ No newline at end of file diff --git a/cpp/src/nn/specializations/fused_l2_knn.cu b/cpp/src/nn/specializations/fused_l2_knn.cu new file mode 100644 index 0000000000..26aa7069e9 --- /dev/null +++ b/cpp/src/nn/specializations/fused_l2_knn.cu @@ -0,0 +1,80 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +namespace raft { +namespace spatial { +namespace knn { +namespace detail { + +template void fusedL2Knn(size_t D, + long* out_inds, + float* out_dists, + const float* index, + const float* query, + size_t n_index_rows, + size_t n_query_rows, + int k, + bool rowMajorIndex, + bool rowMajorQuery, + cudaStream_t stream, + raft::distance::DistanceType metric); + +template void fusedL2Knn(size_t D, + long* out_inds, + float* out_dists, + const float* index, + const float* query, + size_t n_index_rows, + size_t n_query_rows, + int k, + bool rowMajorIndex, + bool rowMajorQuery, + cudaStream_t stream, + raft::distance::DistanceType metric); + +template void fusedL2Knn(size_t D, + int* out_inds, + float* out_dists, + const float* index, + const float* query, + size_t n_index_rows, + size_t n_query_rows, + int k, + bool rowMajorIndex, + bool rowMajorQuery, + cudaStream_t stream, + raft::distance::DistanceType metric); + +template void fusedL2Knn(size_t D, + int* out_inds, + float* out_dists, + const float* index, + const float* query, + size_t n_index_rows, + size_t n_query_rows, + int k, + bool rowMajorIndex, + bool rowMajorQuery, + cudaStream_t stream, + raft::distance::DistanceType metric); + +}; // namespace detail +}; // namespace knn +}; // namespace spatial +}; // namespace raft diff --git a/cpp/src/nn/specializations/knn.cu b/cpp/src/nn/specializations/knn.cu new file mode 100644 index 0000000000..8973cfbb02 --- /dev/null +++ b/cpp/src/nn/specializations/knn.cu @@ -0,0 +1,56 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +namespace raft { +namespace spatial { +namespace knn { + +template void brute_force_knn(raft::handle_t const& handle, + std::vector& input, + std::vector& sizes, + int D, + float* search_items, + int n, + long* res_I, + float* res_D, + int k, + bool rowMajorIndex, + bool rowMajorQuery, + std::vector* translations, + distance::DistanceType metric, + float metric_arg); + +template void brute_force_knn(raft::handle_t const& handle, + std::vector& input, + std::vector& sizes, + unsigned int D, + float* search_items, + unsigned int n, + long* res_I, + float* res_D, + unsigned int k, + bool rowMajorIndex, + bool rowMajorQuery, + std::vector* translations, + distance::DistanceType metric, + float metric_arg); + +}; // namespace knn +}; // namespace spatial +}; // namespace raft diff --git a/cpp/test/CMakeLists.txt b/cpp/test/CMakeLists.txt index 14052293cf..b270204489 100644 --- a/cpp/test/CMakeLists.txt +++ b/cpp/test/CMakeLists.txt @@ -132,7 +132,8 @@ PRIVATE CUDA::cusparse rmm::rmm cuco::cuco - FAISS::FAISS + raft_distance + raft_nn GTest::gtest GTest::gtest_main Threads::Threads diff --git a/cpp/test/distance/distance_base.cuh b/cpp/test/distance/distance_base.cuh index 8d150a4a87..751e895552 100644 --- a/cpp/test/distance/distance_base.cuh +++ b/cpp/test/distance/distance_base.cuh @@ -18,6 +18,7 @@ #include #include #include +#include #include #include "../test_utils.h" @@ -395,12 +396,8 @@ void distanceLauncher(DataType* x, bool isRowMajor, DataType metric_arg = 2.0f) { - auto fin_op = [dist2, threshold] __device__(DataType d_val, int g_d_idx) { - dist2[g_d_idx] = (d_val < threshold) ? 0.f : d_val; - return d_val; - }; raft::distance::distance( - x, y, dist, m, n, k, workspace, worksize, fin_op, stream, isRowMajor, metric_arg); + x, y, dist, m, n, k, workspace, worksize, stream, isRowMajor, metric_arg); } template diff --git a/cpp/test/sparse/knn_graph.cu b/cpp/test/sparse/knn_graph.cu index 9af5e9103b..e7e5854186 100644 --- a/cpp/test/sparse/knn_graph.cu +++ b/cpp/test/sparse/knn_graph.cu @@ -23,6 +23,7 @@ #include #include +#include #include diff --git a/cpp/test/spatial/ball_cover.cu b/cpp/test/spatial/ball_cover.cu index 73c0f87fdd..21d9aaf71f 100644 --- a/cpp/test/spatial/ball_cover.cu +++ b/cpp/test/spatial/ball_cover.cu @@ -18,6 +18,7 @@ #include #include #include +#include #include #include "../test_utils.h" #include "spatial_data.h" diff --git a/cpp/test/spatial/knn.cu b/cpp/test/spatial/knn.cu index 2fb9bd2ca5..5681f66e25 100644 --- a/cpp/test/spatial/knn.cu +++ b/cpp/test/spatial/knn.cu @@ -17,7 +17,9 @@ #include "../test_utils.h" #include + #include +#include #include diff --git a/cpp/test/spatial/selection.cu b/cpp/test/spatial/selection.cu index 5069b4f256..4409f893a8 100644 --- a/cpp/test/spatial/selection.cu +++ b/cpp/test/spatial/selection.cu @@ -21,6 +21,7 @@ #include #include +#include namespace raft { namespace spatial {