Add ray tracing benchmark

benchmarks/CMakeLists.txt

@@ -11,4 +11,5 @@ add_subdirectory(union_find)
 
 if (ARBORX_ENABLE_MPI)
   add_subdirectory(distributed_tree_driver)
+  add_subdirectory(ray_tracing)
 endif()

benchmarks/ray_tracing/CMakeLists.txt

@@ -0,0 +1,5 @@
+add_executable(ArborX_Benchmark_RayTracing.exe ray_tracing.cpp)
+target_link_libraries(ArborX_Benchmark_RayTracing.exe ArborX::ArborX)
+add_test(NAME ArborX_Benchmark_RayTracing COMMAND ${MPIEXEC_EXECUTABLE}
+  ${MPIEXEC_NUMPROC_FLAG} ${MPIEXEC_MAX_NUMPROCS} ${MPIEXEC_PREFLAGS} ./ArborX_Benchmark_RayTracing.exe ${MPIEXEC_POSTFLAGS})
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/rays_cam-0-0_test.csv ${CMAKE_CURRENT_BINARY_DIR}/rays_cam-0-0_test.csv COPYONLY)

benchmarks/ray_tracing/ray_tracing.cpp

@@ -0,0 +1,168 @@
+/****************************************************************************
+ * Copyright (c) 2023 by the ArborX authors                                 *
+ * All rights reserved.                                                     *
+ *                                                                          *
+ * This file is part of the ArborX library. ArborX is                       *
+ * distributed under a BSD 3-clause license. For the licensing terms see    *
+ * the LICENSE file in the top-level directory.                             *
+ *                                                                          *
+ * SPDX-License-Identifier: BSD-3-Clause                                    *
+ ****************************************************************************/
+
+#include <ArborX_DistributedTree.hpp>
+#include <ArborX_Predicates.hpp>
+#include <ArborX_Version.hpp>
+
+#include <Kokkos_Core.hpp>
+
+#include <fstream>
+
+#include <mpi.h>
+
+using ExecutionSpace = Kokkos::DefaultExecutionSpace;
+using MemorySpace = ExecutionSpace::memory_space;
+
+constexpr int n_rays = 307200;
+
+template <typename DataType>
+void read_line(DataType &data, std::string &line)
+{
+  std::size_t pos = 0;
+  std::size_t last_pos = 0;
+  std::size_t line_length = line.length();
+  unsigned int i = 0;
+  while (last_pos < line_length + 1)
+  {
+    pos = line.find_first_of(',', last_pos);
+    // If no comma was found then we read until the end of the file
+    if (pos == std::string::npos)
+    {
+      pos = line_length;
+    }
+
+    if (pos != last_pos)
+    {
+      char *end = line.data() + pos;
+      data[i] = std::strtod(line.data() + last_pos, &end);
+
+      ++i;
+    }
+
+    last_pos = pos + 1;
+  }
+}
+
+template <typename MemorySpace>
+Kokkos::View<ArborX::Experimental::Ray[n_rays], MemorySpace> read_data()
+{
+  std::string filename = "rays_cam-0-0_test.csv";
+  Kokkos::View<ArborX::Experimental::Ray[n_rays], MemorySpace> rays("rays");
+  auto rays_host = Kokkos::create_mirror_view(rays);
+  Kokkos::View<double[n_rays], MemorySpace> temperatures("temperatures");
+
+  std::ifstream file(filename);
+  std::string line;
+  //  First line is the origin of the rays
+  std::getline(file, line);
+  ArborX::Point point;
+  read_line(point, line);
+  // Counter to keep track of the current line
+  int line_number = 0;
+  while (std::getline(file, line))
+  {
+    ArborX::Experimental::Vector direction;
+    read_line(direction, line);
+
+    ArborX::Experimental::Ray ray{point, direction};
+    rays_host(line_number) = ray;
+    ++line_number;
+  }
+  Kokkos::deep_copy(rays, rays_host);
+
+  return rays;
+}
+
+int main(int argc, char *argv[])
+{
+  MPI_Init(&argc, &argv);
+
+  MPI_Comm const comm = MPI_COMM_WORLD;
+  int comm_rank;
+  MPI_Comm_rank(comm, &comm_rank);
+  if (comm_rank == 0)
+  {
+    std::cout << "ArborX version: " << ArborX::version() << std::endl;
+    std::cout << "ArborX hash   : " << ArborX::gitCommitHash() << std::endl;
+    std::cout << "Kokkos version: " << KokkosExt::version() << std::endl;
+  }
+
+  Kokkos::initialize(argc, argv);
+
+  {
+    // Read the experimental data
+    Kokkos::View<ArborX::Experimental::Ray[n_rays], MemorySpace> rays =
+        read_data<MemorySpace>();
+
+    // Build the queries
+    Kokkos::View<ArborX::Nearest<ArborX::Experimental::Ray> *, MemorySpace>
+        queries(Kokkos::view_alloc("queries", Kokkos::WithoutInitializing),
+                rays.size());
+    Kokkos::parallel_for(
+        "Ray::Build_queries",
+        Kokkos::RangePolicy<ExecutionSpace>(0, rays.extent(0)),
+        KOKKOS_LAMBDA(int i) {
+          queries(i) = ArborX::nearest<ArborX::Experimental::Ray>(rays(i), 1);
+        });
+
+    // Create the mesh
+    int comm_size;
+    MPI_Comm_size(comm, &comm_size);
+    constexpr double dim_x = 0.4;
+    constexpr double dim_y = 0.4;
+    constexpr double dim_z = 0.2;
+    // Each processor has one million cells
+    constexpr int n_x = 100;
+    constexpr int n_y = 100;
+    constexpr int n_z = 100;
+    float delta_x = dim_x;
+    float delta_y = dim_y / comm_size;
+    float delta_z = dim_z;
+    float offset_y = comm_rank * delta_y;
+    Kokkos::View<ArborX::Box[n_x * n_y * n_z], MemorySpace> bounding_boxes(
+        Kokkos::view_alloc("bounding_boxes", Kokkos::WithoutInitializing));
+    Kokkos::parallel_for(
+        "build_mesh",
+        Kokkos::MDRangePolicy<Kokkos::Rank<3>>({0, 0, 0}, {n_x, n_y, n_z}),
+        KOKKOS_LAMBDA(int i, int j, int k) {
+          bounding_boxes(i + j * n_x + k * n_x * n_y) = ArborX::Box(
+              {i * delta_x / n_x, offset_y + j * delta_y / n_y,
+               k * delta_z / n_z},
+              {(i + 1) * delta_x / n_x, offset_y + (j + 1) * delta_y / n_y,
+               (k + 1) * delta_z / n_z});
+        });
+
+    // Wait for the processors to be done with the setup phase
+    Kokkos::fence();
+    MPI_Barrier(comm);
+
+    // Build the distributed tree
+    Kokkos::Profiling::pushRegion("Benchmark::Ray::Setup");
+    ArborX::DistributedTree<MemorySpace> distributed_tree(
+        comm, ExecutionSpace{}, bounding_boxes);
+    Kokkos::Profiling::popRegion();
+
+    // Query
+    Kokkos::View<int *, MemorySpace> offsets("offsets", 0);
+    Kokkos::View<ArborX::PairIndexRank *, MemorySpace> indices_ranks(
+        "indices_ranks", 0);
+    Kokkos::Profiling::pushRegion("Benchmark::Ray::Query");
+    distributed_tree.query(ExecutionSpace{}, queries, indices_ranks, offsets);
+    Kokkos::Profiling::popRegion();
+  }
+
+  Kokkos::finalize();
+
+  MPI_Finalize();
+
+  return 0;
+}