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;
     16: using HostMirrorMemorySpace = Kokkos::DualView<PetscScalar *>::host_mirror_space::memory_space;

     18: /* Define a macro if DefaultMemorySpace and HostMirrorMemorySpace are the same */
     19: #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) || defined(KOKKOS_ENABLE_IMPL_CUDA_UNIFIED_MEMORY) || defined(KOKKOS_IMPL_HIP_UNIFIED_MEMORY)
     20:   #define KOKKOS_ENABLE_UNIFIED_MEMORY
     21: #endif

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    127: // Sync a Kokkos::DualView<Type *> to HostMirrorMemorySpace in execution space <exec>
    128: // If <MemorySpace> is HostMirrorMemorySpace, fence the exec so that the data on host is immediately available.
    129: template <typename Type>
    130: static PetscErrorCode KokkosDualViewSyncHost(Kokkos::DualView<Type *> &v_dual, const Kokkos::DefaultExecutionSpace &exec)
    131: {
    132:   size_t bytes = v_dual.extent(0) * sizeof(Type);

    134:   PetscFunctionBegin;
    135:   if (v_dual.need_sync_host()) {
    136:     PetscCallCXX(v_dual.sync_host(exec));
    137:     if (!std::is_same_v<DefaultMemorySpace, HostMirrorMemorySpace>) PetscCall(PetscLogGpuToCpu(bytes));
    138:   }
    139:   // even if v_d and v_h share the same memory (as on AMD MI300A) and thus we don't need to sync_host,
    140:   // we still need to fence the execution space as v_d might being populated by some async kernel,
    141:   // and we need to finish it.
    142:   PetscCallCXX(exec.fence());
    143:   PetscFunctionReturn(PETSC_SUCCESS);
    144: }

    146: template <typename Type>
    147: static PetscErrorCode KokkosDualViewSyncDevice(Kokkos::DualView<Type *> &v_dual, const Kokkos::DefaultExecutionSpace &exec)
    148: {
    149:   size_t bytes = v_dual.extent(0) * sizeof(Type);

    151:   PetscFunctionBegin;
    152:   if (v_dual.need_sync_device()) {
    153:     PetscCallCXX(v_dual.sync_device(exec));
    154:     if (!std::is_same_v<DefaultMemorySpace, HostMirrorMemorySpace>) PetscCall(PetscLogCpuToGpu(bytes));
    155:   }
    156:   PetscFunctionReturn(PETSC_SUCCESS);
    157: }