Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Removing code that explicitly compares equality of rmm memory resources #2047

Merged
merged 14 commits into from
Dec 9, 2023
Prev Previous commit
Next Next commit
More fixes
  • Loading branch information
cjnolet committed Dec 8, 2023
commit 8bdfe330467786fa1281aa0d2c5d1c2255651b2e
6 changes: 4 additions & 2 deletions cpp/include/raft/neighbors/detail/ivf_pq_build.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@
#include <raft/core/resource/device_memory_resource.hpp>
#include <rmm/cuda_stream_view.hpp>
#include <rmm/device_uvector.hpp>
#include <rmm/mr/device/managed_memory_resource.hpp>

#include <thrust/extrema.h>
#include <thrust/scan.h>
Expand Down Expand Up @@ -1755,6 +1756,7 @@ auto build(raft::resources const& handle,
size_t n_rows_train = n_rows / trainset_ratio;

auto* device_memory = resource::get_workspace_resource(handle);
rmm::mr::managed_memory_resource managed_memory_upstream;

// Besides just sampling, we transform the input dataset into floats to make it easier
// to use gemm operations from cublas.
Expand Down Expand Up @@ -1856,7 +1858,7 @@ auto build(raft::resources const& handle,
trainset.data(),
labels.data(),
params.kmeans_n_iters,
device_memory);
&managed_memory_upstream);
break;
case codebook_gen::PER_CLUSTER:
train_per_cluster(handle,
Expand All @@ -1865,7 +1867,7 @@ auto build(raft::resources const& handle,
trainset.data(),
labels.data(),
params.kmeans_n_iters,
device_memory);
&managed_memory_upstream);
break;
default: RAFT_FAIL("Unreachable code");
}
Expand Down
Loading