Actual source code: vecimpl_kokkos.hpp

  1: #ifndef __VECIMPL_KOKKOS_HPP

  4: /* types used by all petsc kokkos code */

  6: #include <petscvec_kokkos.hpp>
  7: #include <Kokkos_Core.hpp>
  8: #include <Kokkos_DualView.hpp>
  9: #include <Kokkos_OffsetView.hpp>

 11: using DefaultExecutionSpace = Kokkos::DefaultExecutionSpace;
 12: using DefaultMemorySpace    = Kokkos::DefaultExecutionSpace::memory_space;

 14: /* Define a macro if DefaultMemorySpace is HostSpace */
 15: #if defined(KOKKOS_ENABLE_DEFAULT_DEVICE_TYPE_SERIAL) || defined(KOKKOS_ENABLE_DEFAULT_DEVICE_TYPE_OPENMP) || defined(KOKKOS_ENABLE_DEFAULT_DEVICE_TYPE_THREADS) || defined(KOKKOS_ENABLE_DEFAULT_DEVICE_TYPE_HPX)
 16:   #define KOKKOS_ENABLE_DEFAULT_DEVICE_TYPE_HOST
 17: #endif

 19: /* 1 to 4D PetscScalar Kokkos Views */
 20: template <class MemorySpace>
 21: using PetscScalarKokkosViewType = Kokkos::View<PetscScalar *, MemorySpace>;
 22: template <class MemorySpace>
 23: using PetscScalarKokkosView1DType = Kokkos::View<PetscScalar *, MemorySpace>;
 24: template <class MemorySpace>
 25: using PetscScalarKokkosView2DType = Kokkos::View<PetscScalar **, Kokkos::LayoutRight, MemorySpace>;
 26: template <class MemorySpace>
 27: using PetscScalarKokkosView3DType = Kokkos::View<PetscScalar ***, Kokkos::LayoutRight, MemorySpace>;
 28: template <class MemorySpace>
 29: using PetscScalarKokkosView4DType = Kokkos::View<PetscScalar ****, Kokkos::LayoutRight, MemorySpace>;

 31: template <class MemorySpace>
 32: using ConstPetscScalarKokkosViewType = Kokkos::View<const PetscScalar *, MemorySpace>;
 33: template <class MemorySpace>
 34: using ConstPetscScalarKokkosView1DType = Kokkos::View<const PetscScalar *, MemorySpace>;
 35: template <class MemorySpace>
 36: using ConstPetscScalarKokkosView2DType = Kokkos::View<const PetscScalar **, Kokkos::LayoutRight, MemorySpace>;
 37: template <class MemorySpace>
 38: using ConstPetscScalarKokkosView3DType = Kokkos::View<const PetscScalar ***, Kokkos::LayoutRight, MemorySpace>;
 39: template <class MemorySpace>
 40: using ConstPetscScalarKokkosView4DType = Kokkos::View<const PetscScalar ****, Kokkos::LayoutRight, MemorySpace>;

 42: /* 1 to 4D PetscScalar Kokkos OffsetViews */
 43: template <class MemorySpace>
 44: using PetscScalarKokkosOffsetViewType = Kokkos::Experimental::OffsetView<PetscScalar *, MemorySpace>;
 45: template <class MemorySpace>
 46: using PetscScalarKokkosOffsetView1DType = Kokkos::Experimental::OffsetView<PetscScalar *, MemorySpace>;
 47: template <class MemorySpace>
 48: using PetscScalarKokkosOffsetView2DType = Kokkos::Experimental::OffsetView<PetscScalar **, Kokkos::LayoutRight, MemorySpace>;
 49: template <class MemorySpace>
 50: using PetscScalarKokkosOffsetView3DType = Kokkos::Experimental::OffsetView<PetscScalar ***, Kokkos::LayoutRight, MemorySpace>;
 51: template <class MemorySpace>
 52: using PetscScalarKokkosOffsetView4DType = Kokkos::Experimental::OffsetView<PetscScalar ****, Kokkos::LayoutRight, MemorySpace>;

 54: template <class MemorySpace>
 55: using ConstPetscScalarKokkosOffsetViewType = Kokkos::Experimental::OffsetView<const PetscScalar *, MemorySpace>;
 56: template <class MemorySpace>
 57: using ConstPetscScalarKokkosOffsetView1DType = Kokkos::Experimental::OffsetView<const PetscScalar *, MemorySpace>;
 58: template <class MemorySpace>
 59: using ConstPetscScalarKokkosOffsetView2DType = Kokkos::Experimental::OffsetView<const PetscScalar **, Kokkos::LayoutRight, MemorySpace>;
 60: template <class MemorySpace>
 61: using ConstPetscScalarKokkosOffsetView3DType = Kokkos::Experimental::OffsetView<const PetscScalar ***, Kokkos::LayoutRight, MemorySpace>;
 62: template <class MemorySpace>
 63: using ConstPetscScalarKokkosOffsetView4DType = Kokkos::Experimental::OffsetView<const PetscScalar ****, Kokkos::LayoutRight, MemorySpace>;

 65: using PetscScalarKokkosDualView = Kokkos::DualView<PetscScalar *>;

 67: /* Shortcut types for Views in the default space and host space */
 68: using PetscScalarKokkosView   = PetscScalarKokkosViewType<DefaultMemorySpace>;
 69: using PetscScalarKokkosView1D = PetscScalarKokkosView1DType<DefaultMemorySpace>;
 70: using PetscScalarKokkosView2D = PetscScalarKokkosView2DType<DefaultMemorySpace>;
 71: using PetscScalarKokkosView3D = PetscScalarKokkosView3DType<DefaultMemorySpace>;
 72: using PetscScalarKokkosView4D = PetscScalarKokkosView4DType<DefaultMemorySpace>;

 74: using PetscScalarKokkosViewHost   = PetscScalarKokkosViewType<Kokkos::HostSpace>;
 75: using PetscScalarKokkosView1DHost = PetscScalarKokkosView1DType<Kokkos::HostSpace>;
 76: using PetscScalarKokkosView2DHost = PetscScalarKokkosView2DType<Kokkos::HostSpace>;
 77: using PetscScalarKokkosView3DHost = PetscScalarKokkosView3DType<Kokkos::HostSpace>;
 78: using PetscScalarKokkosView4DHost = PetscScalarKokkosView4DType<Kokkos::HostSpace>;

 80: using ConstPetscScalarKokkosView   = ConstPetscScalarKokkosViewType<DefaultMemorySpace>;
 81: using ConstPetscScalarKokkosView1D = ConstPetscScalarKokkosView1DType<DefaultMemorySpace>;
 82: using ConstPetscScalarKokkosView2D = ConstPetscScalarKokkosView2DType<DefaultMemorySpace>;
 83: using ConstPetscScalarKokkosView3D = ConstPetscScalarKokkosView3DType<DefaultMemorySpace>;
 84: using ConstPetscScalarKokkosView4D = ConstPetscScalarKokkosView4DType<DefaultMemorySpace>;

 86: using ConstPetscScalarKokkosViewHost   = ConstPetscScalarKokkosViewType<Kokkos::HostSpace>;
 87: using ConstPetscScalarKokkosView1DHost = ConstPetscScalarKokkosView1DType<Kokkos::HostSpace>;
 88: using ConstPetscScalarKokkosView2DHost = ConstPetscScalarKokkosView2DType<Kokkos::HostSpace>;
 89: using ConstPetscScalarKokkosView3DHost = ConstPetscScalarKokkosView3DType<Kokkos::HostSpace>;
 90: using ConstPetscScalarKokkosView4DHost = ConstPetscScalarKokkosView4DType<Kokkos::HostSpace>;

 92: /* Shortcut types for OffsetViews in the default space and host space */
 93: using PetscScalarKokkosOffsetView   = PetscScalarKokkosOffsetViewType<DefaultMemorySpace>;
 94: using PetscScalarKokkosOffsetView1D = PetscScalarKokkosOffsetView1DType<DefaultMemorySpace>;
 95: using PetscScalarKokkosOffsetView2D = PetscScalarKokkosOffsetView2DType<DefaultMemorySpace>;
 96: using PetscScalarKokkosOffsetView3D = PetscScalarKokkosOffsetView3DType<DefaultMemorySpace>;
 97: using PetscScalarKokkosOffsetView4D = PetscScalarKokkosOffsetView4DType<DefaultMemorySpace>;

 99: using PetscScalarKokkosOffsetViewHost   = PetscScalarKokkosOffsetViewType<Kokkos::HostSpace>;
100: using PetscScalarKokkosOffsetView1DHost = PetscScalarKokkosOffsetView1DType<Kokkos::HostSpace>;
101: using PetscScalarKokkosOffsetView2DHost = PetscScalarKokkosOffsetView2DType<Kokkos::HostSpace>;
102: using PetscScalarKokkosOffsetView3DHost = PetscScalarKokkosOffsetView3DType<Kokkos::HostSpace>;
103: using PetscScalarKokkosOffsetView4DHost = PetscScalarKokkosOffsetView4DType<Kokkos::HostSpace>;

105: using ConstPetscScalarKokkosOffsetView   = ConstPetscScalarKokkosOffsetViewType<DefaultMemorySpace>;
106: using ConstPetscScalarKokkosOffsetView1D = ConstPetscScalarKokkosOffsetView1DType<DefaultMemorySpace>;
107: using ConstPetscScalarKokkosOffsetView2D = ConstPetscScalarKokkosOffsetView2DType<DefaultMemorySpace>;
108: using ConstPetscScalarKokkosOffsetView3D = ConstPetscScalarKokkosOffsetView3DType<DefaultMemorySpace>;
109: using ConstPetscScalarKokkosOffsetView4D = ConstPetscScalarKokkosOffsetView4DType<DefaultMemorySpace>;

111: using ConstPetscScalarKokkosOffsetViewHost   = ConstPetscScalarKokkosOffsetViewType<Kokkos::HostSpace>;
112: using ConstPetscScalarKokkosOffsetView1DHost = ConstPetscScalarKokkosOffsetView1DType<Kokkos::HostSpace>;
113: using ConstPetscScalarKokkosOffsetView2DHost = ConstPetscScalarKokkosOffsetView2DType<Kokkos::HostSpace>;
114: using ConstPetscScalarKokkosOffsetView3DHost = ConstPetscScalarKokkosOffsetView3DType<Kokkos::HostSpace>;
115: using ConstPetscScalarKokkosOffsetView4DHost = ConstPetscScalarKokkosOffsetView4DType<Kokkos::HostSpace>;

117: using PetscIntKokkosView       = Kokkos::View<PetscInt *, DefaultMemorySpace>;
118: using PetscIntKokkosViewHost   = Kokkos::View<PetscInt *, Kokkos::HostSpace>;
119: using PetscIntKokkosDualView   = Kokkos::DualView<PetscInt *>;
120: using PetscCountKokkosView     = Kokkos::View<PetscCount *, DefaultMemorySpace>;
121: using PetscCountKokkosViewHost = Kokkos::View<PetscCount *, Kokkos::HostSpace>;
122: #endif