Actual source code: kokkosimpl.hpp

      1: #pragma once

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

      5: #include <petscsystypes.h>
      6: #include <Kokkos_Core.hpp>
      7: #include <Kokkos_DualView.hpp>
      8: #include <Kokkos_OffsetView.hpp>

     10: // the pool is defined in veckok.kokkos.cxx as it is currently only used there
     11: PETSC_SINGLE_LIBRARY_INTERN PetscScalar *PetscScalarPool;
     12: PETSC_SINGLE_LIBRARY_INTERN PetscInt     PetscScalarPoolSize;

     14: using DefaultExecutionSpace = Kokkos::DefaultExecutionSpace;
     15: using DefaultMemorySpace    = Kokkos::DefaultExecutionSpace::memory_space;

     17: /* Define a macro if DefaultMemorySpace is HostSpace */
     18: #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)
     19:   #define KOKKOS_ENABLE_DEFAULT_DEVICE_TYPE_HOST
     20: #endif

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

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

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

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

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

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

     77: using PetscScalarKokkosViewHost   = PetscScalarKokkosViewType<Kokkos::HostSpace>;
     78: using PetscScalarKokkosView1DHost = PetscScalarKokkosView1DType<Kokkos::HostSpace>;
     79: using PetscScalarKokkosView2DHost = PetscScalarKokkosView2DType<Kokkos::HostSpace>;
     80: using PetscScalarKokkosView3DHost = PetscScalarKokkosView3DType<Kokkos::HostSpace>;
     81: using PetscScalarKokkosView4DHost = PetscScalarKokkosView4DType<Kokkos::HostSpace>;

     83: using ConstPetscScalarKokkosView   = ConstPetscScalarKokkosViewType<DefaultMemorySpace>;
     84: using ConstPetscScalarKokkosView1D = ConstPetscScalarKokkosView1DType<DefaultMemorySpace>;
     85: using ConstPetscScalarKokkosView2D = ConstPetscScalarKokkosView2DType<DefaultMemorySpace>;
     86: using ConstPetscScalarKokkosView3D = ConstPetscScalarKokkosView3DType<DefaultMemorySpace>;
     87: using ConstPetscScalarKokkosView4D = ConstPetscScalarKokkosView4DType<DefaultMemorySpace>;

     89: using ConstPetscScalarKokkosViewHost   = ConstPetscScalarKokkosViewType<Kokkos::HostSpace>;
     90: using ConstPetscScalarKokkosView1DHost = ConstPetscScalarKokkosView1DType<Kokkos::HostSpace>;
     91: using ConstPetscScalarKokkosView2DHost = ConstPetscScalarKokkosView2DType<Kokkos::HostSpace>;
     92: using ConstPetscScalarKokkosView3DHost = ConstPetscScalarKokkosView3DType<Kokkos::HostSpace>;
     93: using ConstPetscScalarKokkosView4DHost = ConstPetscScalarKokkosView4DType<Kokkos::HostSpace>;

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

    102: using PetscScalarKokkosOffsetViewHost   = PetscScalarKokkosOffsetViewType<Kokkos::HostSpace>;
    103: using PetscScalarKokkosOffsetView1DHost = PetscScalarKokkosOffsetView1DType<Kokkos::HostSpace>;
    104: using PetscScalarKokkosOffsetView2DHost = PetscScalarKokkosOffsetView2DType<Kokkos::HostSpace>;
    105: using PetscScalarKokkosOffsetView3DHost = PetscScalarKokkosOffsetView3DType<Kokkos::HostSpace>;
    106: using PetscScalarKokkosOffsetView4DHost = PetscScalarKokkosOffsetView4DType<Kokkos::HostSpace>;

    108: using ConstPetscScalarKokkosOffsetView   = ConstPetscScalarKokkosOffsetViewType<DefaultMemorySpace>;
    109: using ConstPetscScalarKokkosOffsetView1D = ConstPetscScalarKokkosOffsetView1DType<DefaultMemorySpace>;
    110: using ConstPetscScalarKokkosOffsetView2D = ConstPetscScalarKokkosOffsetView2DType<DefaultMemorySpace>;
    111: using ConstPetscScalarKokkosOffsetView3D = ConstPetscScalarKokkosOffsetView3DType<DefaultMemorySpace>;
    112: using ConstPetscScalarKokkosOffsetView4D = ConstPetscScalarKokkosOffsetView4DType<DefaultMemorySpace>;

    114: using ConstPetscScalarKokkosOffsetViewHost   = ConstPetscScalarKokkosOffsetViewType<Kokkos::HostSpace>;
    115: using ConstPetscScalarKokkosOffsetView1DHost = ConstPetscScalarKokkosOffsetView1DType<Kokkos::HostSpace>;
    116: using ConstPetscScalarKokkosOffsetView2DHost = ConstPetscScalarKokkosOffsetView2DType<Kokkos::HostSpace>;
    117: using ConstPetscScalarKokkosOffsetView3DHost = ConstPetscScalarKokkosOffsetView3DType<Kokkos::HostSpace>;
    118: using ConstPetscScalarKokkosOffsetView4DHost = ConstPetscScalarKokkosOffsetView4DType<Kokkos::HostSpace>;

    120: using PetscIntKokkosView       = Kokkos::View<PetscInt *, DefaultMemorySpace>;
    121: using PetscIntKokkosViewHost   = Kokkos::View<PetscInt *, Kokkos::HostSpace>;
    122: using PetscIntKokkosDualView   = Kokkos::DualView<PetscInt *>;
    123: using PetscCountKokkosView     = Kokkos::View<PetscCount *, DefaultMemorySpace>;
    124: using PetscCountKokkosViewHost = Kokkos::View<PetscCount *, Kokkos::HostSpace>;