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 to in execution space
128: // If is HostMirrorMemorySpace, fence the exec so that the data on host is immediately available.
129: template <class MemorySpace, typename Type>
130: static PetscErrorCode KokkosDualViewSync(Kokkos::DualView<Type *> &v_dual, const Kokkos::DefaultExecutionSpace &exec)
131: {
132: size_t bytes = v_dual.extent(0) * sizeof(Type);
134: PetscFunctionBegin;
135: if (std::is_same_v<MemorySpace, HostMirrorMemorySpace>) {
136: if (v_dual.need_sync_host()) {
137: PetscCallCXX(v_dual.sync_host(exec));
138: PetscCall(PetscLogGpuToCpu(bytes));
139: }
140: // even if v_d and v_h share the same memory (as on AMD MI300A) and thus we don't need to sync_host,
141: // we still need to fence the execution space as v_d might being populated by some async kernel,
142: // and we need to finish it.
143: PetscCallCXX(exec.fence());
144: } else {
145: if (v_dual.need_sync_device()) {
146: PetscCallCXX(v_dual.sync_device(exec));
147: PetscCall(PetscLogCpuToGpu(bytes));
148: }
149: }
150: PetscFunctionReturn(PETSC_SUCCESS);
151: }