Merge pull request #12089 from NexGenAnalytics/nga-fy23-pr-candidate-3

Zoltan2: Refactor of BasisVectorAdapter to use Kokkos data structures
trilinos · Sep 7, 2023 · 8c6ce31 · 8c6ce31
2 parents 5dff89a + 83c26bb
commit 8c6ce31
Show file tree

Hide file tree

Showing 5 changed files with 461 additions and 128 deletions.
diff --git a/packages/zoltan2/core/src/input/Zoltan2_Adapter.hpp b/packages/zoltan2/core/src/input/Zoltan2_Adapter.hpp
@@ -336,23 +336,27 @@ template <typename User>
 
 };
 
-template <typename User> class AdapterWithCoords : public BaseAdapter<User>
+template <typename User>
+class AdapterWithCoords : public BaseAdapter<User>
 {
+public:
   using scalar_t = typename BaseAdapter<User>::scalar_t;
   using device_t = typename BaseAdapter<User>::node_t::device_type;
   using host_t = typename Kokkos::HostSpace::memory_space;
-  template <typename space>
-  using coords_t = Kokkos::View<scalar_t **, Kokkos::LayoutLeft, space>;
+
+  // Coordinates in MJ are LayoutLeft since Tpetra Multivector gives LayoutLeft
+  using CoordsDeviceView = Kokkos::View<scalar_t **, Kokkos::LayoutLeft, device_t>;
+  using CoordsHostView = typename CoordsDeviceView::HostMirror;
 
 public:
   virtual void getCoordinatesView(const scalar_t *&coords, int &stride,
                                   int coordDim) const = 0;
-  virtual void getCoordinatesKokkosView(coords_t<device_t> &elements) const = 0;
+  virtual void getCoordinatesKokkosView(CoordsDeviceView &elements) const = 0;
 
-  virtual void getCoordinatesHostView(coords_t<host_t> &) const {
+  virtual void getCoordinatesHostView(CoordsHostView &) const {
     Z2_THROW_NOT_IMPLEMENTED
   }
-  virtual void getCoordinatesDeviceView(coords_t<device_t> &elements) const {
+  virtual void getCoordinatesDeviceView(CoordsDeviceView &elements) const {
     Z2_THROW_NOT_IMPLEMENTED
   }
 };