From a277388e3e5c5a905b0a5006f9c8591a7efbf05f Mon Sep 17 00:00:00 2001 From: "Artem M. Chirkin" <9253178+achirkin@users.noreply.github.com> Date: Sat, 18 Feb 2023 15:01:35 +0100 Subject: [PATCH] matrix::select_k specializations (#1268) Add specializations to `matrix::detail::select_k`, which has a few use-cases across raft. Authors: - Artem M. Chirkin (https://github.com/achirkin) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/raft/pull/1268 --- cpp/CMakeLists.txt | 4 ++ .../specializations/detail/select_k.cuh | 47 +++++++++++++++++++ .../raft/neighbors/detail/ivf_pq_build.cuh | 3 +- .../raft/neighbors/specializations/ivf_pq.cuh | 1 + .../raft_internal/matrix/select_k.cuh | 6 +++ .../raft_internal/neighbors/naive_knn.cuh | 4 ++ .../detail/select_k_float_uint32_t.cu | 35 ++++++++++++++ .../detail/select_k_float_uint64_t.cu | 35 ++++++++++++++ .../detail/select_k_half_uint32_t.cu | 35 ++++++++++++++ .../detail/select_k_half_uint64_t.cu | 35 ++++++++++++++ cpp/test/neighbors/ann_ivf_flat.cuh | 3 ++ cpp/test/neighbors/selection.cu | 3 ++ 12 files changed, 210 insertions(+), 1 deletion(-) create mode 100644 cpp/include/raft/matrix/specializations/detail/select_k.cuh create mode 100644 cpp/src/distance/matrix/specializations/detail/select_k_float_uint32_t.cu create mode 100644 cpp/src/distance/matrix/specializations/detail/select_k_float_uint64_t.cu create mode 100644 cpp/src/distance/matrix/specializations/detail/select_k_half_uint32_t.cu create mode 100644 cpp/src/distance/matrix/specializations/detail/select_k_half_uint64_t.cu diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index b020b8421f..74a4d31b67 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -358,6 +358,10 @@ if(RAFT_COMPILE_DIST_LIBRARY) src/distance/distance/specializations/fused_l2_nn_double_int64.cu src/distance/distance/specializations/fused_l2_nn_float_int.cu src/distance/distance/specializations/fused_l2_nn_float_int64.cu + src/distance/matrix/specializations/detail/select_k_float_uint32_t.cu + src/distance/matrix/specializations/detail/select_k_float_uint64_t.cu + src/distance/matrix/specializations/detail/select_k_half_uint32_t.cu + src/distance/matrix/specializations/detail/select_k_half_uint64_t.cu src/distance/neighbors/ivfpq_build.cu src/distance/neighbors/ivfpq_deserialize.cu src/distance/neighbors/ivfpq_serialize.cu diff --git a/cpp/include/raft/matrix/specializations/detail/select_k.cuh b/cpp/include/raft/matrix/specializations/detail/select_k.cuh new file mode 100644 index 0000000000..3326162812 --- /dev/null +++ b/cpp/include/raft/matrix/specializations/detail/select_k.cuh @@ -0,0 +1,47 @@ +/* + * Copyright (c) 2023, 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::matrix::detail { + +#define RAFT_INST(T, IdxT) \ + extern template void select_k(const T*, \ + const IdxT*, \ + size_t, \ + size_t, \ + int, \ + T*, \ + IdxT*, \ + bool, \ + rmm::cuda_stream_view, \ + rmm::mr::device_memory_resource*); + +// Commonly used types +RAFT_INST(float, uint64_t); +RAFT_INST(half, uint64_t); + +// These instances are used in the ivf_pq::search parameterized by the internal_distance_dtype +RAFT_INST(float, uint32_t); +RAFT_INST(half, uint32_t); + +#undef RAFT_INST + +} // namespace raft::matrix::detail diff --git a/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh b/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh index bc334b28f1..bf3014568a 100644 --- a/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh +++ b/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh @@ -34,7 +34,8 @@ #include #include #include -#include +#include +#include #include #include #include diff --git a/cpp/include/raft/neighbors/specializations/ivf_pq.cuh b/cpp/include/raft/neighbors/specializations/ivf_pq.cuh index a7afc9f361..3ff99fb4da 100644 --- a/cpp/include/raft/neighbors/specializations/ivf_pq.cuh +++ b/cpp/include/raft/neighbors/specializations/ivf_pq.cuh @@ -16,6 +16,7 @@ #pragma once +#include #include #include diff --git a/cpp/internal/raft_internal/matrix/select_k.cuh b/cpp/internal/raft_internal/matrix/select_k.cuh index 205149b821..398c6c8d1b 100644 --- a/cpp/internal/raft_internal/matrix/select_k.cuh +++ b/cpp/internal/raft_internal/matrix/select_k.cuh @@ -14,10 +14,16 @@ * limitations under the License. */ +#pragma once + #include #include #include +#if defined RAFT_DISTANCE_COMPILED +#include +#endif + #include namespace raft::matrix::select { diff --git a/cpp/internal/raft_internal/neighbors/naive_knn.cuh b/cpp/internal/raft_internal/neighbors/naive_knn.cuh index 3ad055272b..942c096e58 100644 --- a/cpp/internal/raft_internal/neighbors/naive_knn.cuh +++ b/cpp/internal/raft_internal/neighbors/naive_knn.cuh @@ -21,6 +21,10 @@ #include #include +#if defined RAFT_DISTANCE_COMPILED +#include +#endif + #include #include #include diff --git a/cpp/src/distance/matrix/specializations/detail/select_k_float_uint32_t.cu b/cpp/src/distance/matrix/specializations/detail/select_k_float_uint32_t.cu new file mode 100644 index 0000000000..30b7bbbd3e --- /dev/null +++ b/cpp/src/distance/matrix/specializations/detail/select_k_float_uint32_t.cu @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2023, 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::matrix::detail { + +#define RAFT_INST(T, IdxT) \ + template void select_k(const T*, \ + const IdxT*, \ + size_t, \ + size_t, \ + int, \ + T*, \ + IdxT*, \ + bool, \ + rmm::cuda_stream_view, \ + rmm::mr::device_memory_resource*); + +RAFT_INST(float, uint32_t); + +} // namespace raft::matrix::detail diff --git a/cpp/src/distance/matrix/specializations/detail/select_k_float_uint64_t.cu b/cpp/src/distance/matrix/specializations/detail/select_k_float_uint64_t.cu new file mode 100644 index 0000000000..c449bf03e4 --- /dev/null +++ b/cpp/src/distance/matrix/specializations/detail/select_k_float_uint64_t.cu @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2023, 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::matrix::detail { + +#define RAFT_INST(T, IdxT) \ + template void select_k(const T*, \ + const IdxT*, \ + size_t, \ + size_t, \ + int, \ + T*, \ + IdxT*, \ + bool, \ + rmm::cuda_stream_view, \ + rmm::mr::device_memory_resource*); + +RAFT_INST(float, uint64_t); + +} // namespace raft::matrix::detail diff --git a/cpp/src/distance/matrix/specializations/detail/select_k_half_uint32_t.cu b/cpp/src/distance/matrix/specializations/detail/select_k_half_uint32_t.cu new file mode 100644 index 0000000000..129fc44c81 --- /dev/null +++ b/cpp/src/distance/matrix/specializations/detail/select_k_half_uint32_t.cu @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2023, 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::matrix::detail { + +#define RAFT_INST(T, IdxT) \ + template void select_k(const T*, \ + const IdxT*, \ + size_t, \ + size_t, \ + int, \ + T*, \ + IdxT*, \ + bool, \ + rmm::cuda_stream_view, \ + rmm::mr::device_memory_resource*); + +RAFT_INST(half, uint32_t); + +} // namespace raft::matrix::detail diff --git a/cpp/src/distance/matrix/specializations/detail/select_k_half_uint64_t.cu b/cpp/src/distance/matrix/specializations/detail/select_k_half_uint64_t.cu new file mode 100644 index 0000000000..b08dc8641c --- /dev/null +++ b/cpp/src/distance/matrix/specializations/detail/select_k_half_uint64_t.cu @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2023, 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::matrix::detail { + +#define RAFT_INST(T, IdxT) \ + template void select_k(const T*, \ + const IdxT*, \ + size_t, \ + size_t, \ + int, \ + T*, \ + IdxT*, \ + bool, \ + rmm::cuda_stream_view, \ + rmm::mr::device_memory_resource*); + +RAFT_INST(half, uint64_t); + +} // namespace raft::matrix::detail diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh index c100afb2c4..26b8301cb1 100644 --- a/cpp/test/neighbors/ann_ivf_flat.cuh +++ b/cpp/test/neighbors/ann_ivf_flat.cuh @@ -36,6 +36,9 @@ #include +#if defined RAFT_DISTANCE_COMPILED +#include +#endif #if defined RAFT_DISTANCE_COMPILED && defined RAFT_NN_COMPILED #include #endif diff --git a/cpp/test/neighbors/selection.cu b/cpp/test/neighbors/selection.cu index 61a6345e5e..872cb1873e 100644 --- a/cpp/test/neighbors/selection.cu +++ b/cpp/test/neighbors/selection.cu @@ -24,6 +24,9 @@ #include #include +#if defined RAFT_DISTANCE_COMPILED +#include +#endif #if defined RAFT_NN_COMPILED #include #endif