1 #pragma once
2
3 /* types used by all PETSc Kokkos code */
4
5 #include <petscsystypes.h>
6 #include <Kokkos_Core.hpp>
7 #include <Kokkos_DualView.hpp>
8 #include <Kokkos_OffsetView.hpp>
9
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;
13
14 using DefaultExecutionSpace = Kokkos::DefaultExecutionSpace;
15 using DefaultMemorySpace = Kokkos::DefaultExecutionSpace::memory_space;
16 using HostMirrorMemorySpace = Kokkos::DualView<PetscScalar *>::host_mirror_space::memory_space;
17
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
22
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>;
34
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>;
45
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>;
57
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>;
68
69 using PetscScalarKokkosDualView = Kokkos::DualView<PetscScalar *>;
70
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>;
77
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>;
83
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>;
89
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>;
95
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>;
102
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>;
108
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>;
114
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>;
120
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>;
126
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>
KokkosDualViewSyncHost(Kokkos::DualView<Type * > & v_dual,const Kokkos::DefaultExecutionSpace & exec)130 static PetscErrorCode KokkosDualViewSyncHost(Kokkos::DualView<Type *> &v_dual, const Kokkos::DefaultExecutionSpace &exec)
131 {
132 size_t bytes = v_dual.extent(0) * sizeof(Type);
133
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 }
145
146 template <typename Type>
KokkosDualViewSyncDevice(Kokkos::DualView<Type * > & v_dual,const Kokkos::DefaultExecutionSpace & exec)147 static PetscErrorCode KokkosDualViewSyncDevice(Kokkos::DualView<Type *> &v_dual, const Kokkos::DefaultExecutionSpace &exec)
148 {
149 size_t bytes = v_dual.extent(0) * sizeof(Type);
150
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 }
158