Merge pull request #348 from Fuad-HH/debug-stdarray-gpu

Fix std::array Bug for GPU kernels - For Conservative Transfer

Author: Fuad-HH

src/pcms/localization/point_search.cpp

@@ -65,8 +65,8 @@ LO GetOwningElementId(Omega_h::Mesh& mesh, int mesh_dim, int entity_dim,
 KOKKOS_INLINE_FUNCTION
 AABBox<2> triangle_bbox(const Omega_h::Matrix<2, 3>& coords)
 {
-  std::array<Real, 2> max{coords(0, 0), coords(1, 0)};
-  std::array<Real, 2> min{coords(0, 0), coords(1, 0)};
+  Kokkos::Array<Real, 2> max{coords(0, 0), coords(1, 0)};
+  Kokkos::Array<Real, 2> min{coords(0, 0), coords(1, 0)};
   for (int i = 1; i < 3; ++i) {
     max[0] = std::fmax(max[0], coords(0, i));
     max[1] = std::fmax(max[1], coords(1, i));
@@ -185,7 +185,7 @@ template <int dim>
   // each dimension has a pair of opposing "walls"
   // 2D: { [left, right], [top, bottom] } -> { left, right, top, bottom }
   // 3D: { [left, right], [top, bottom], [front, back] } -> { left, ..., back }
-  std::array<Real, dim * 2ul> bbox_walls{};
+  Kokkos::Array<Real, dim * 2ul> bbox_walls{};
   for (int i = 0; i < dim; i++) {
     bbox_walls[i * 2] = bbox.center[i] - bbox.half_width[i];
     bbox_walls[i * 2 + 1] = bbox.center[i] + bbox.half_width[i];
@@ -751,8 +751,8 @@ GridPointSearch3D::GridPointSearch3D(Omega_h::Mesh& mesh, LO Nx, LO Ny, LO Nz,
   auto mesh_bbox = Omega_h::get_bounding_box<3>(&mesh);
   auto grid_h = Kokkos::create_mirror_view(grid_);
 
-  std::array<Real, DIM> edge_lengths{};
-  std::array<Real, DIM> bot_left{};
+  Kokkos::Array<Real, DIM> edge_lengths{};
+  Kokkos::Array<Real, DIM> bot_left{};
 
   for (int i = 0; i < DIM; ++i) {
     edge_lengths[i] = mesh_bbox.max[i] - mesh_bbox.min[i];

src/pcms/pythonapi/bind_uniform_grid_field_layout.cpp

@@ -1,6 +1,7 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 #include <pybind11/numpy.h>
+#include "numpy_array_transform.h"
 #include "pcms/field/layout/uniform_grid.h"
 #include "pcms/utility/uniform_grid.h"
 

src/pcms/pythonapi/numpy_array_transform.h

@@ -7,6 +7,48 @@
 
 namespace py = pybind11;
 
+// Mirror pybind11's built-in std::array type caster for Kokkos::Array so that
+// def_readwrite / def_property members work with Python sequences.
+namespace pybind11::detail
+{
+template <typename Type, size_t Size>
+struct type_caster<Kokkos::Array<Type, Size>>
+{
+  using ArrayType = Kokkos::Array<Type, Size>;
+  using value_conv = make_caster<Type>;
+
+  bool load(handle src, bool convert)
+  {
+    if (!isinstance<sequence>(src))
+      return false;
+    auto seq = reinterpret_borrow<sequence>(src);
+    if (seq.size() != Size)
+      return false;
+    size_t i = 0;
+    for (auto it : seq) {
+      value_conv conv;
+      if (!conv.load(it, convert))
+        return false;
+      value[i++] = cast_op<Type>(conv);
+    }
+    return true;
+  }
+
+  template <typename T>
+  static handle cast(T&& src, return_value_policy policy, handle parent)
+  {
+    list l(Size);
+    size_t i = 0;
+    for (auto& v : src)
+      l[i++] = make_caster<Type>::cast(v, policy, parent);
+    return l.release();
+  }
+
+  PYBIND11_TYPE_CASTER(ArrayType, const_name("List[") +
+                                    make_caster<Type>::name + const_name("]"));
+};
+} // namespace pybind11::detail
+
 namespace pcms
 {
 

src/pcms/utility/arrays.h

@@ -70,6 +70,12 @@ struct arr_trait<std::array<T, N>>
   using type = T;
 };
 
+template <typename T, size_t N>
+struct arr_trait<Kokkos::Array<T, N>>
+{
+  using type = T;
+};
+
 template <typename T>
 using element_type_t = typename arr_trait<T, HasValueType<T>::value>::type;
 

src/pcms/utility/uniform_grid.h

@@ -4,23 +4,26 @@
 #include "Omega_h_vector.hpp"
 #include "Omega_h_bbox.hpp"
 #include "Omega_h_mesh.hpp"
-#include <numeric>
+#include <Kokkos_Array.hpp>
+#include <array>
 namespace pcms
 {
 
 template <unsigned dim>
 struct UniformGrid
 {
   // Make private?
-  std::array<Real, dim> edge_length;
-  std::array<Real, dim> bot_left;
-  std::array<LO, dim> divisions;
+  Kokkos::Array<Real, dim> edge_length;
+  Kokkos::Array<Real, dim> bot_left;
+  Kokkos::Array<LO, dim> divisions;
 
 public:
   [[nodiscard]] LO GetNumCells() const
   {
-    return std::accumulate(divisions.begin(), divisions.end(), 1,
-                           std::multiplies<LO>{});
+    LO total = 1;
+    for (std::size_t i = 0; i < dim; ++i)
+      total *= divisions[i];
+    return total;
   }
   /// return the grid cell ID that the input point is inside or closest to if
   /// the point lies outside
@@ -30,13 +33,13 @@ struct UniformGrid
   [[nodiscard]] KOKKOS_INLINE_FUNCTION LO
   ClosestCellID(const Omega_h::Vector<dim>& point) const
   {
-    std::array<Real, dim> distance_within_grid;
+    Kokkos::Array<Real, dim> distance_within_grid;
 
     for (size_t i = 0; i < dim; ++i) {
       distance_within_grid[i] = point[i] - bot_left[i];
     }
 
-    std::array<LO, dim> indexes;
+    Kokkos::Array<LO, dim> indexes;
 
     for (auto& index : indexes) {
       index = -1;
@@ -86,14 +89,14 @@ struct UniformGrid
     return true;
   }
 
-  [[nodiscard]] KOKKOS_INLINE_FUNCTION std::array<LO, dim> GetDimensionedIndex(
-    LO idx) const
+  [[nodiscard]] KOKKOS_INLINE_FUNCTION Kokkos::Array<LO, dim>
+  GetDimensionedIndex(LO idx) const
   {
     LO stride = 1;
-    for (std::size_t i = 0; i < divisions.size() - 1; ++i) {
+    for (std::size_t i = 0; i < dim - 1; ++i) {
       stride *= divisions[i];
     }
-    std::array<LO, dim> result;
+    Kokkos::Array<LO, dim> result;
 
     for (size_t i = 0; i < dim; ++i) {
       result[i] = idx / stride;
@@ -105,7 +108,7 @@ struct UniformGrid
   }
 
   [[nodiscard]] KOKKOS_INLINE_FUNCTION LO
-  GetCellIndex(std::array<LO, dim> dimensionedIndex) const
+  GetCellIndex(Kokkos::Array<LO, dim> dimensionedIndex) const
   {
     // note that the indexes refer to row/columns which have the opposite order
     // of the coordinates i.e. x,y
@@ -124,9 +127,9 @@ struct UniformGrid
 
 private:
   template <typename T, std::size_t N>
-  KOKKOS_INLINE_FUNCTION static void reverse(std::array<T, N>& arr)
+  KOKKOS_INLINE_FUNCTION static void reverse(Kokkos::Array<T, N>& arr)
   {
-    for (size_t i = 0, j = arr.size() - 1; i < j; ++i, --j) {
+    for (size_t i = 0, j = N - 1; i < j; ++i, --j) {
       auto temp = arr[i];
       arr[i] = arr[j];
       arr[j] = temp;
@@ -158,16 +161,18 @@ UniformGrid<dim> CreateUniformGridFromMesh(Omega_h::Mesh& mesh,
   auto bbox = Omega_h::get_bounding_box<dim>(&mesh);
 
   // Calculate edge lengths and bottom-left corner
-  std::array<Real, dim> edge_length;
-  std::array<Real, dim> bot_left;
+  Kokkos::Array<Real, dim> edge_length;
+  Kokkos::Array<Real, dim> bot_left;
+  Kokkos::Array<LO, dim> divs;
 
   for (unsigned i = 0; i < dim; ++i) {
     bot_left[i] = bbox.min[i];
     edge_length[i] = bbox.max[i] - bbox.min[i];
+    divs[i] = divisions[i];
   }
 
   return UniformGrid<dim>{
-    .edge_length = edge_length, .bot_left = bot_left, .divisions = divisions};
+    .edge_length = edge_length, .bot_left = bot_left, .divisions = divs};
 }
 
 /**
@@ -187,7 +192,8 @@ UniformGrid<dim> CreateUniformGridFromMesh(Omega_h::Mesh& mesh,
                                            LO cells_per_dim)
 {
   std::array<LO, dim> divisions;
-  divisions.fill(cells_per_dim);
+  for (unsigned i = 0; i < dim; ++i)
+    divisions[i] = cells_per_dim;
   return CreateUniformGridFromMesh<dim>(mesh, divisions);
 }