-
Notifications
You must be signed in to change notification settings - Fork 197
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #2232 from rapidsai/branch-24.04
Forward-merge branch-24.04 into branch-24.06 [skip ci]
- Loading branch information
Showing
41 changed files
with
3,067 additions
and
179 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
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,112 @@ | ||
| /* | ||
| * Copyright (c) 2024, 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 <common/benchmark.hpp> | ||
|
|
||
| #include <raft/core/device_mdarray.hpp> | ||
| #include <raft/core/device_resources.hpp> | ||
| #include <raft/core/host_mdarray.hpp> | ||
| #include <raft/core/operators.hpp> | ||
| #include <raft/random/permute.cuh> | ||
| #include <raft/random/rng.cuh> | ||
| #include <raft/random/sample_without_replacement.cuh> | ||
| #include <raft/spatial/knn/detail/ann_utils.cuh> | ||
| #include <raft/util/cudart_utils.hpp> | ||
|
|
||
| #include <rmm/device_scalar.hpp> | ||
| #include <rmm/mr/device/pool_memory_resource.hpp> | ||
|
|
||
| #include <cub/cub.cuh> | ||
|
|
||
| namespace raft::bench::random { | ||
|
|
||
| struct sample_inputs { | ||
| int n_samples; | ||
| int n_train; | ||
| int method; | ||
| }; // struct sample_inputs | ||
|
|
||
| inline auto operator<<(std::ostream& os, const sample_inputs& p) -> std::ostream& | ||
| { | ||
| os << p.n_samples << "#" << p.n_train << "#" << p.method; | ||
| return os; | ||
| } | ||
|
|
||
| // Sample with replacement. We use this as a baseline. | ||
| template <typename IdxT> | ||
| auto bernoulli_subsample(raft::resources const& res, IdxT n_samples, IdxT n_subsamples, int seed) | ||
| -> raft::device_vector<IdxT, IdxT> | ||
| { | ||
| RAFT_EXPECTS(n_subsamples <= n_samples, "Cannot have more training samples than dataset vectors"); | ||
|
|
||
| auto indices = raft::make_device_vector<IdxT, IdxT>(res, n_subsamples); | ||
| raft::random::RngState state(123456ULL); | ||
| raft::random::uniformInt( | ||
| res, state, indices.data_handle(), n_subsamples, IdxT(0), IdxT(n_samples)); | ||
| return indices; | ||
| } | ||
|
|
||
| template <typename T> | ||
| struct sample : public fixture { | ||
| sample(const sample_inputs& p) | ||
| : params(p), | ||
| old_mr(rmm::mr::get_current_device_resource()), | ||
| pool_mr(rmm::mr::get_current_device_resource(), 2 * GiB), | ||
| in(make_device_vector<T, int64_t>(res, p.n_samples)), | ||
| out(make_device_vector<T, int64_t>(res, p.n_train)) | ||
| { | ||
| rmm::mr::set_current_device_resource(&pool_mr); | ||
| raft::random::RngState r(123456ULL); | ||
| } | ||
|
|
||
| ~sample() { rmm::mr::set_current_device_resource(old_mr); } | ||
| void run_benchmark(::benchmark::State& state) override | ||
| { | ||
| std::ostringstream label_stream; | ||
| label_stream << params; | ||
| state.SetLabel(label_stream.str()); | ||
|
|
||
| raft::random::RngState r(123456ULL); | ||
| loop_on_state(state, [this, &r]() { | ||
| if (params.method == 1) { | ||
| this->out = | ||
| bernoulli_subsample<T>(this->res, this->params.n_samples, this->params.n_train, 137); | ||
| } else if (params.method == 2) { | ||
| this->out = raft::random::excess_subsample<T, int64_t>( | ||
| this->res, r, this->params.n_samples, this->params.n_train); | ||
| } | ||
| }); | ||
| } | ||
|
|
||
| private: | ||
| float GiB = 1073741824.0f; | ||
| raft::device_resources res; | ||
| rmm::mr::device_memory_resource* old_mr; | ||
| rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr; | ||
| sample_inputs params; | ||
| raft::device_vector<T, int64_t> out, in; | ||
| }; // struct sample | ||
|
|
||
| const std::vector<sample_inputs> input_vecs = {{100000000, 10000000, 1}, | ||
| {100000000, 50000000, 1}, | ||
| {100000000, 100000000, 1}, | ||
| {100000000, 10000000, 2}, | ||
| {100000000, 50000000, 2}, | ||
| {100000000, 100000000, 2}}; | ||
|
|
||
| RAFT_BENCH_REGISTER(sample<int64_t>, "", input_vecs); | ||
|
|
||
| } // namespace raft::bench::random |
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,97 @@ | ||
| /* | ||
| * Copyright (c) 2024, 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 <raft/core/kvp.hpp> // raft::KeyValuePair | ||
| #include <raft/core/operators.hpp> // raft::identity_op | ||
| #include <raft/distance/detail/distance_ops/l2_exp.cuh> // ops::l2_exp_distance_op | ||
| #include <raft/distance/detail/fused_distance_nn/cutlass_base.cuh> | ||
| #include <raft/distance/detail/fused_distance_nn/fused_cosine_nn.cuh> | ||
| #include <raft/distance/detail/fused_distance_nn/fused_l2_nn.cuh> | ||
| #include <raft/distance/detail/fused_distance_nn/helper_structs.cuh> | ||
| #include <raft/distance/detail/fused_distance_nn/simt_kernel.cuh> | ||
| #include <raft/distance/detail/pairwise_distance_base.cuh> // PairwiseDistances | ||
| #include <raft/distance/distance_types.hpp> | ||
| #include <raft/linalg/contractions.cuh> // Policy | ||
| #include <raft/util/arch.cuh> // raft::util::arch::SM_* | ||
| #include <raft/util/cuda_utils.cuh> // raft::ceildiv, raft::shfl | ||
|
|
||
| #include <cstddef> // size_t | ||
| #include <limits> // std::numeric_limits | ||
|
|
||
| namespace raft { | ||
| namespace distance { | ||
|
|
||
| namespace detail { | ||
|
|
||
| template <typename DataT, | ||
| typename OutT, | ||
| typename IdxT, | ||
| typename Policy, | ||
| typename ReduceOpT, | ||
| typename KVPReduceOpT> | ||
| void fusedDistanceNNImpl(OutT* min, | ||
| const DataT* x, | ||
| const DataT* y, | ||
| const DataT* xn, | ||
| const DataT* yn, | ||
| IdxT m, | ||
| IdxT n, | ||
| IdxT k, | ||
| int* workspace, | ||
| ReduceOpT redOp, | ||
| KVPReduceOpT pairRedOp, | ||
| bool sqrt, | ||
| bool initOutBuffer, | ||
| bool isRowMajor, | ||
| raft::distance::DistanceType metric, | ||
| float metric_arg, | ||
| cudaStream_t stream) | ||
| { | ||
| // The kernel policy is determined by fusedDistanceNN. | ||
| typedef Policy P; | ||
|
|
||
| dim3 blk(P::Nthreads); | ||
| auto nblks = raft::ceildiv<int>(m, P::Nthreads); | ||
| constexpr auto maxVal = std::numeric_limits<DataT>::max(); | ||
| typedef KeyValuePair<IdxT, DataT> KVPair; | ||
|
|
||
| RAFT_CUDA_TRY(cudaMemsetAsync(workspace, 0, sizeof(int) * m, stream)); | ||
| if (initOutBuffer) { | ||
| initKernel<DataT, OutT, IdxT, ReduceOpT> | ||
| <<<nblks, P::Nthreads, 0, stream>>>(min, m, maxVal, redOp); | ||
| RAFT_CUDA_TRY(cudaGetLastError()); | ||
| } | ||
|
|
||
| switch (metric) { | ||
| case raft::distance::DistanceType::CosineExpanded: | ||
| fusedCosineNN<DataT, OutT, IdxT, P, ReduceOpT, KVPReduceOpT>( | ||
| min, x, y, xn, yn, m, n, k, workspace, redOp, pairRedOp, sqrt, stream); | ||
| break; | ||
| case raft::distance::DistanceType::L2SqrtExpanded: | ||
| case raft::distance::DistanceType::L2Expanded: | ||
| // initOutBuffer is take care by fusedDistanceNNImpl() so we set it false to fusedL2NNImpl. | ||
| fusedL2NNImpl<DataT, OutT, IdxT, P, ReduceOpT, KVPReduceOpT>( | ||
| min, x, y, xn, yn, m, n, k, workspace, redOp, pairRedOp, sqrt, false, stream); | ||
| break; | ||
| default: assert("only cosine/l2 metric is supported with fusedDistanceNN\n"); break; | ||
| } | ||
| } | ||
|
|
||
| } // namespace detail | ||
| } // namespace distance | ||
| } // namespace raft |
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
Oops, something went wrong.