Moving more prims benchmarks to RAFT

cpp/bench/CMakeLists.txt

@@ -18,7 +18,18 @@ set(RAFT_CPP_BENCH_TARGET "bench_raft")
 
 # (please keep the filenames in alphabetical order)
 add_executable(${RAFT_CPP_BENCH_TARGET}
+  bench/distance/distance_cosine.cu
+  bench/distance/distance_exp_l2.cu
+  bench/distance/distance_l1.cu
+  bench/distance/distance_unexp_l2.cu
+  bench/linalg/add.cu
+  bench/linalg/map_then_reduce.cu
+  bench/linalg/matrix_vector_op.cu
   bench/linalg/reduce.cu
+  bench/random/make_blobs.cu
+  bench/random/permute.cu
+  bench/random/rng.cu
+  bench/spatial/fused_l2_nn.cu
   bench/spatial/selection.cu
   bench/main.cpp
 )
@@ -47,6 +58,8 @@ target_include_directories(${RAFT_CPP_BENCH_TARGET}
 target_link_libraries(${RAFT_CPP_BENCH_TARGET}
   PRIVATE
     raft::raft
+    raft::distance
+    raft::nn
     faiss::faiss
     benchmark::benchmark
     $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>

cpp/bench/distance/distance_common.cuh

@@ -0,0 +1,95 @@
+/*
+ * Copyright (c) 2022, 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/cudart_utils.h>
+#include <raft/distance/distance.hpp>
+#if defined RAFT_DISTANCE_COMPILED
+#include <raft/distance/specializations.hpp>
+#endif
+#include <rmm/device_uvector.hpp>
+
+namespace raft::bench::distance {
+
+struct distance_inputs {
+  int m, n, k;
+  bool isRowMajor;
+};  // struct distance_inputs
+
+template <typename T, raft::distance::DistanceType DType>
+struct distance : public fixture {
+  distance(const distance_inputs& p)
+    : params(p),
+      x(p.m * p.k, stream),
+      y(p.n * p.k, stream),
+      out(p.m * p.n, stream),
+      workspace(0, stream)
+  {
+    RAFT_CUDA_TRY(cudaMemsetAsync(x.data(), 0, x.size() * sizeof(T), stream));
+    RAFT_CUDA_TRY(cudaMemsetAsync(y.data(), 0, y.size() * sizeof(T), stream));
+    RAFT_CUDA_TRY(cudaMemsetAsync(out.data(), 0, out.size() * sizeof(T), stream));
+    worksize = raft::distance::getWorkspaceSize<DType, T, T, T>(
+      x.data(), y.data(), params.m, params.n, params.k);
+    workspace.resize(worksize, stream);
+  }
+
+  void run_benchmark(::benchmark::State& state) override
+  {
+    loop_on_state(state, [this]() {
+      raft::distance::distance<DType, T, T, T>(x.data(),
+                                               y.data(),
+                                               out.data(),
+                                               params.m,
+                                               params.n,
+                                               params.k,
+                                               (void*)workspace.data(),
+                                               worksize,
+                                               stream,
+                                               params.isRowMajor);
+    });
+  }
+
+ private:
+  distance_inputs params;
+  rmm::device_uvector<T> x, y, out;
+  rmm::device_uvector<char> workspace;
+  size_t worksize;
+};  // struct Distance
+
+const std::vector<distance_inputs> dist_input_vecs{
+  {32, 16384, 16384, true},    {64, 16384, 16384, true},    {128, 16384, 16384, true},
+  {256, 16384, 16384, true},   {512, 16384, 16384, true},   {1024, 16384, 16384, true},
+  {16384, 32, 16384, true},    {16384, 64, 16384, true},    {16384, 128, 16384, true},
+  {16384, 256, 16384, true},   {16384, 512, 16384, true},   {16384, 1024, 16384, true},
+  {16384, 16384, 32, true},    {16384, 16384, 64, true},    {16384, 16384, 128, true},
+  {16384, 16384, 256, true},   {16384, 16384, 512, true},   {16384, 16384, 1024, true},
+  {16384, 16384, 16384, true}, {32, 16384, 16384, false},   {64, 16384, 16384, false},
+  {128, 16384, 16384, false},  {256, 16384, 16384, false},  {512, 16384, 16384, false},
+  {1024, 16384, 16384, false}, {16384, 32, 16384, false},   {16384, 64, 16384, false},
+  {16384, 128, 16384, false},  {16384, 256, 16384, false},  {16384, 512, 16384, false},
+  {16384, 1024, 16384, false}, {16384, 16384, 32, false},   {16384, 16384, 64, false},
+  {16384, 16384, 128, false},  {16384, 16384, 256, false},  {16384, 16384, 512, false},
+  {16384, 16384, 1024, false}, {16384, 16384, 16384, false}
+
+};
+
+#define DIST_BENCH_REGISTER(Name, Metric)            \
+  using Name##F = distance<float, Metric>;           \
+  RAFT_BENCH_REGISTER(Name##F, "", dist_input_vecs); \
+  using Name##D = distance<double, Metric>;          \
+  RAFT_BENCH_REGISTER(Name##D, "", dist_input_vecs);
+
+}  // namespace raft::bench::distance

cpp/bench/distance/distance_cosine.cu

@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) 2022, 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 "distance_common.cuh"
+
+namespace raft::bench::distance {
+
+DIST_BENCH_REGISTER(DistanceCosine, raft::distance::DistanceType::CosineExpanded);
+
+}  // namespace raft::bench::distance

cpp/bench/distance/distance_exp_l2.cu

@@ -0,0 +1,24 @@
+/*
+ * Copyright (c) 2022, 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 "distance_common.cuh"
+
+namespace raft::bench::distance {
+
+DIST_BENCH_REGISTER(DistanceL2Sq, raft::distance::DistanceType::L2Expanded);
+DIST_BENCH_REGISTER(DistanceL2Sqrt, raft::distance::DistanceType::L2SqrtExpanded);
+
+}  // namespace raft::bench::distance

cpp/bench/distance/distance_l1.cu

@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) 2022, 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 "distance_common.cuh"
+
+namespace raft::bench::distance {
+
+DIST_BENCH_REGISTER(DistanceL1, raft::distance::DistanceType::L1);
+
+}  // namespace raft::bench::distance

cpp/bench/distance/distance_unexp_l2.cu

@@ -0,0 +1,24 @@
+/*
+ * Copyright (c) 2022, 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 "distance_common.cuh"
+
+namespace raft::bench::distance {
+
+DIST_BENCH_REGISTER(DistanceUnexpL2Sq, raft::distance::DistanceType::L2Unexpanded);
+DIST_BENCH_REGISTER(DistanceUnexpL2Sqrt, raft::distance::DistanceType::L2SqrtUnexpanded);
+
+}  // namespace raft::bench::distance

cpp/bench/linalg/add.cu

@@ -0,0 +1,51 @@
+/*
+ * Copyright (c) 2022, 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/linalg/add.hpp>
+#include <rmm/device_uvector.hpp>
+
+namespace raft::bench::linalg {
+
+struct add_inputs {
+  int len;
+};  // struct add_inputs
+
+template <typename T>
+struct add : public fixture {
+  add(const add_inputs& p) : params(p), ptr0(p.len, stream), ptr1(p.len, stream) {}
+
+  void run_benchmark(::benchmark::State& state) override
+  {
+    loop_on_state(state, [this]() {
+      raft::linalg::add(ptr0.data(), ptr0.data(), ptr1.data(), params.len, stream);
+    });
+  }
+
+ private:
+  add_inputs params;
+  rmm::device_uvector<T> ptr0, ptr1;
+};  // struct add
+
+const std::vector<add_inputs> add_input_vecs{
+  {256 * 1024 * 1024}, {256 * 1024 * 1024 + 2}, {256 * 1024 * 1024 + 1}
+
+};
+
+RAFT_BENCH_REGISTER(add<float>, "", add_input_vecs);
+RAFT_BENCH_REGISTER(add<double>, "", add_input_vecs);
+
+}  // namespace raft::bench::linalg

cpp/bench/linalg/map_then_reduce.cu

@@ -0,0 +1,64 @@
+/*
+ * Copyright (c) 2022, 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/linalg/map_then_reduce.hpp>
+#include <rmm/device_uvector.hpp>
+
+namespace raft::bench::linalg {
+
+struct map_then_reduce_inputs {
+  int len;
+};
+
+template <typename Type>
+struct Identity {
+  HDI Type operator()(Type a) { return a; }
+};
+
+template <typename T>
+struct map_then_reduce : public fixture {
+  map_then_reduce(const map_then_reduce_inputs& p) : params(p), in(p.len, stream), out(1, stream) {}
+
+  void run_benchmark(::benchmark::State& state) override
+  {
+    loop_on_state(state, [this]() {
+      raft::linalg::mapThenSumReduce(out.data(), params.len, Identity<T>(), stream, in.data());
+    });
+  }
+
+ private:
+  map_then_reduce_inputs params;
+  rmm::device_uvector<T> out, in;
+};  // struct MapThenReduce
+
+const std::vector<map_then_reduce_inputs> map_then_reduce_input_vecs{
+  {1024 * 1024},
+  {32 * 1024 * 1024},
+  {1024 * 1024 * 1024},
+  {1024 * 1024 + 2},
+  {32 * 1024 * 1024 + 2},
+  {1024 * 1024 * 1024 + 2},
+  {1024 * 1024 + 1},
+  {32 * 1024 * 1024 + 1},
+  {1024 * 1024 * 1024 + 1},
+
+};
+
+RAFT_BENCH_REGISTER(map_then_reduce<float>, "", map_then_reduce_input_vecs);
+RAFT_BENCH_REGISTER(map_then_reduce<double>, "", map_then_reduce_input_vecs);
+
+}  // namespace raft::bench::linalg