forked from rapidsai/raft
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Split memory_pool, fused_l2_knn, coalesced_reduction, selection_faiss
- Loading branch information
1 parent
95638fd
commit 8a4c4a8
Showing
19 changed files
with
620 additions
and
80 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
74 changes: 74 additions & 0 deletions
74
cpp/include/raft/linalg/detail/coalesced_reduction-ext.cuh
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,74 @@ | ||
| /* | ||
| * Copyright (c) 2022-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 "coalesced_reduction-types.cuh" | ||
| #include <raft/core/operators.hpp> | ||
|
|
||
| // Include inline definition as well. We cannot possibly cover all | ||
| // instantiations in this file. | ||
| #include "coalesced_reduction-inl.cuh" | ||
|
|
||
| #define instantiate_raft_linalg_detail_coalescedReduction( \ | ||
| InType, OutType, IdxType, MainLambda, ReduceLambda, FinalLambda) \ | ||
| extern template void raft::linalg::detail::coalescedReduction(OutType* dots, \ | ||
| const InType* data, \ | ||
| IdxType D, \ | ||
| IdxType N, \ | ||
| OutType init, \ | ||
| cudaStream_t stream, \ | ||
| bool inplace, \ | ||
| MainLambda main_op, \ | ||
| ReduceLambda reduce_op, \ | ||
| FinalLambda final_op) | ||
|
|
||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| double, double, int, raft::identity_op, raft::min_op, raft::identity_op); | ||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| double, double, int, raft::sq_op, raft::add_op, raft::identity_op); | ||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| double, double, int, raft::sq_op, raft::add_op, raft::sqrt_op); | ||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| double, double, int, raft::abs_op, raft::add_op, raft::identity_op); | ||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| double, double, int, raft::abs_op, raft::max_op, raft::identity_op); | ||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| float, float, size_t, raft::abs_op, raft::add_op, raft::sqrt_op); | ||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| float, float, int, raft::abs_op, raft::add_op, raft::identity_op); | ||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| float, float, int, raft::identity_op, raft::add_op, raft::identity_op); | ||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| float, float, int, raft::identity_op, raft::min_op, raft::identity_op); | ||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| float, float, int, raft::sq_op, raft::add_op, raft::identity_op); | ||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| float, float, int, raft::sq_op, raft::add_op, raft::sqrt_op); | ||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| float, float, long, raft::sq_op, raft::add_op, raft::identity_op); | ||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| float, float, size_t, raft::identity_op, raft::add_op, raft::identity_op); | ||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| float, float, size_t, raft::sq_op, raft::add_op, raft::identity_op); | ||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| float, float, size_t, raft::abs_op, raft::max_op, raft::sqrt_op); | ||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| float, float, size_t, raft::sq_op, raft::add_op, raft::sqrt_op); | ||
| instantiate_raft_linalg_detail_coalescedReduction( | ||
| float, float, unsigned int, raft::sq_op, raft::add_op, raft::identity_op); | ||
|
|
||
| #undef instantiate_raft_linalg_detail_coalescedReduction |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
34 changes: 34 additions & 0 deletions
34
cpp/include/raft/linalg/detail/coalesced_reduction-types.cuh
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,34 @@ | ||
| /* | ||
| * 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 | ||
|
|
||
| namespace raft::linalg::detail { | ||
|
|
||
| template <int warpSize, int rpb> | ||
| struct ReductionThinPolicy { | ||
| static constexpr int LogicalWarpSize = warpSize; | ||
| static constexpr int RowsPerBlock = rpb; | ||
| static constexpr int ThreadsPerBlock = LogicalWarpSize * RowsPerBlock; | ||
| }; | ||
|
|
||
| template <int tpb, int bpr> | ||
| struct ReductionThickPolicy { | ||
| static constexpr int ThreadsPerBlock = tpb; | ||
| static constexpr int BlocksPerRow = bpr; | ||
| static constexpr int BlockStride = tpb * bpr; | ||
| }; | ||
|
|
||
| } // namespace raft::linalg::detail |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,25 @@ | ||
| /* | ||
| * Copyright (c) 2022-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 | ||
|
|
||
| #if defined(RAFT_COMPILED) && defined(RAFT_EXPLICIT_INSTANTIATE) | ||
| // Too many lambdas and complicated types to instantiate everything.. | ||
| #include "coalesced_reduction-ext.cuh" | ||
| #include "coalesced_reduction-inl.cuh" | ||
| #else | ||
| #include "coalesced_reduction-inl.cuh" | ||
| #endif |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,70 @@ | ||
| /* | ||
| * Copyright (c) 2020-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 <cstddef> // size_t | ||
| #include <cstdint> // uint32_t | ||
| #include <raft/util/raft_explicit.hpp> // RAFT_EXPLICIT | ||
|
|
||
| #if defined(RAFT_EXPLICIT_INSTANTIATE) | ||
|
|
||
| namespace raft::neighbors::detail { | ||
| /** | ||
| * @brief Select the k-nearest neighbors from dense | ||
| * distance and index matrices. | ||
| * | ||
| * @param[in] inK partitioned knn distance matrix | ||
| * @param[in] inV partitioned knn index matrix | ||
| * @param[in] n_rows number of rows in distance and index matrices | ||
| * @param[in] n_cols number of columns in distance and index matrices | ||
| * @param[out] outK merged knn distance matrix | ||
| * @param[out] outV merged knn index matrix | ||
| * @param[in] select_min whether to select the min or the max distances | ||
| * @param[in] k number of neighbors per partition (also number of merged neighbors) | ||
| * @param[in] stream CUDA stream to use | ||
| */ | ||
| template <typename payload_t = int, typename key_t = float> | ||
| void select_k(const key_t* inK, | ||
| const payload_t* inV, | ||
| size_t n_rows, | ||
| size_t n_cols, | ||
| key_t* outK, | ||
| payload_t* outV, | ||
| bool select_min, | ||
| int k, | ||
| cudaStream_t stream) RAFT_EXPLICIT; | ||
| }; // namespace raft::neighbors::detail | ||
|
|
||
| #endif // RAFT_EXPLICIT_INSTANTIATE | ||
|
|
||
| #define instantiate_raft_neighbors_detail_select_k(payload_t, key_t) \ | ||
| extern template void raft::neighbors::detail::select_k(const key_t* inK, \ | ||
| const payload_t* inV, \ | ||
| size_t n_rows, \ | ||
| size_t n_cols, \ | ||
| key_t* outK, \ | ||
| payload_t* outV, \ | ||
| bool select_min, \ | ||
| int k, \ | ||
| cudaStream_t stream) | ||
|
|
||
| // @benfred: Not sure if this is correct. Should I not flip float and uint32_t? | ||
| // It seems weird that float is the key and uint32_t is the payload type. | ||
| instantiate_raft_neighbors_detail_select_k(uint32_t, float); | ||
| instantiate_raft_neighbors_detail_select_k(long, float); | ||
|
|
||
| #undef instantiate_raft_neighbors_detail_select_k |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,25 @@ | ||
| /* | ||
| * Copyright (c) 2020-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 | ||
|
|
||
| #if defined(RAFT_COMPILED) | ||
| #include "selection_faiss-ext.cuh" | ||
| #endif | ||
|
|
||
| #if !defined(RAFT_EXPLICIT_INSTANTIATE) | ||
| #include "selection_faiss-inl.cuh" | ||
| #endif |
Oops, something went wrong.