Actual source code: kokkosimpl.hpp
1: #pragma once
3: /* types used by all PETSc Kokkos code */
5: #include <petscsystypes.h>
6: #include <petscerror.h>
7: #include <Kokkos_Core.hpp>
8: #include <Kokkos_DualView.hpp>
9: #include <Kokkos_OffsetView.hpp>
11: // the pool is defined in veckok.kokkos.cxx as it is currently only used there
12: PETSC_SINGLE_LIBRARY_INTERN PetscScalar *PetscScalarPool;
13: PETSC_SINGLE_LIBRARY_INTERN PetscInt PetscScalarPoolSize;
15: using DefaultExecutionSpace = Kokkos::DefaultExecutionSpace;
16: using DefaultMemorySpace = Kokkos::DefaultExecutionSpace::memory_space;
17: using HostMirrorMemorySpace = Kokkos::DualView<PetscScalar *>::host_mirror_space::memory_space;
19: /* Define a macro if DefaultMemorySpace and HostMirrorMemorySpace are the same */
20: #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)
21: #define KOKKOS_ENABLE_UNIFIED_MEMORY
22: #endif
24: /* 1 to 4D PetscScalar Kokkos Views */
25: template <class MemorySpace>
26: using PetscScalarKokkosViewType = Kokkos::View<PetscScalar *, MemorySpace>;
27: template <class MemorySpace>
28: using PetscScalarKokkosView1DType = Kokkos::View<PetscScalar *, MemorySpace>;
29: template <class MemorySpace>
30: using PetscScalarKokkosView2DType = Kokkos::View<PetscScalar **, Kokkos::LayoutRight, MemorySpace>;
31: template <class MemorySpace>
32: using PetscScalarKokkosView3DType = Kokkos::View<PetscScalar ***, Kokkos::LayoutRight, MemorySpace>;
33: template <class MemorySpace>
34: using PetscScalarKokkosView4DType = Kokkos::View<PetscScalar ****, Kokkos::LayoutRight, MemorySpace>;
36: template <class MemorySpace>
37: using ConstPetscScalarKokkosViewType = Kokkos::View<const PetscScalar *, MemorySpace>;
38: template <class MemorySpace>
39: using ConstPetscScalarKokkosView1DType = Kokkos::View<const PetscScalar *, MemorySpace>;
40: template <class MemorySpace>
41: using ConstPetscScalarKokkosView2DType = Kokkos::View<const PetscScalar **, Kokkos::LayoutRight, MemorySpace>;
42: template <class MemorySpace>
43: using ConstPetscScalarKokkosView3DType = Kokkos::View<const PetscScalar ***, Kokkos::LayoutRight, MemorySpace>;
44: template <class MemorySpace>
45: using ConstPetscScalarKokkosView4DType = Kokkos::View<const PetscScalar ****, Kokkos::LayoutRight, MemorySpace>;
47: /* 1 to 4D PetscScalar Kokkos OffsetViews */
48: template <class MemorySpace>
49: using PetscScalarKokkosOffsetViewType = Kokkos::Experimental::OffsetView<PetscScalar *, MemorySpace>;
50: template <class MemorySpace>
51: using PetscScalarKokkosOffsetView1DType = Kokkos::Experimental::OffsetView<PetscScalar *, MemorySpace>;
52: template <class MemorySpace>
53: using PetscScalarKokkosOffsetView2DType = Kokkos::Experimental::OffsetView<PetscScalar **, Kokkos::LayoutRight, MemorySpace>;
54: template <class MemorySpace>
55: using PetscScalarKokkosOffsetView3DType = Kokkos::Experimental::OffsetView<PetscScalar ***, Kokkos::LayoutRight, MemorySpace>;
56: template <class MemorySpace>
57: using PetscScalarKokkosOffsetView4DType = Kokkos::Experimental::OffsetView<PetscScalar ****, Kokkos::LayoutRight, MemorySpace>;
59: template <class MemorySpace>
60: using ConstPetscScalarKokkosOffsetViewType = Kokkos::Experimental::OffsetView<const PetscScalar *, MemorySpace>;
61: template <class MemorySpace>
62: using ConstPetscScalarKokkosOffsetView1DType = Kokkos::Experimental::OffsetView<const PetscScalar *, MemorySpace>;
63: template <class MemorySpace>
64: using ConstPetscScalarKokkosOffsetView2DType = Kokkos::Experimental::OffsetView<const PetscScalar **, Kokkos::LayoutRight, MemorySpace>;
65: template <class MemorySpace>
66: using ConstPetscScalarKokkosOffsetView3DType = Kokkos::Experimental::OffsetView<const PetscScalar ***, Kokkos::LayoutRight, MemorySpace>;
67: template <class MemorySpace>
68: using ConstPetscScalarKokkosOffsetView4DType = Kokkos::Experimental::OffsetView<const PetscScalar ****, Kokkos::LayoutRight, MemorySpace>;
70: using PetscScalarKokkosDualView = Kokkos::DualView<PetscScalar *>;
72: /* Shortcut types for Views in the default space and host space */
73: using PetscScalarKokkosView = PetscScalarKokkosViewType<DefaultMemorySpace>;
74: using PetscScalarKokkosView1D = PetscScalarKokkosView1DType<DefaultMemorySpace>;
75: using PetscScalarKokkosView2D = PetscScalarKokkosView2DType<DefaultMemorySpace>;
76: using PetscScalarKokkosView3D = PetscScalarKokkosView3DType<DefaultMemorySpace>;
77: using PetscScalarKokkosView4D = PetscScalarKokkosView4DType<DefaultMemorySpace>;
79: using PetscScalarKokkosViewHost = PetscScalarKokkosViewType<HostMirrorMemorySpace>;
80: using PetscScalarKokkosView1DHost = PetscScalarKokkosView1DType<HostMirrorMemorySpace>;
81: using PetscScalarKokkosView2DHost = PetscScalarKokkosView2DType<HostMirrorMemorySpace>;
82: using PetscScalarKokkosView3DHost = PetscScalarKokkosView3DType<HostMirrorMemorySpace>;
83: using PetscScalarKokkosView4DHost = PetscScalarKokkosView4DType<HostMirrorMemorySpace>;
85: using ConstPetscScalarKokkosView = ConstPetscScalarKokkosViewType<DefaultMemorySpace>;
86: using ConstPetscScalarKokkosView1D = ConstPetscScalarKokkosView1DType<DefaultMemorySpace>;
87: using ConstPetscScalarKokkosView2D = ConstPetscScalarKokkosView2DType<DefaultMemorySpace>;
88: using ConstPetscScalarKokkosView3D = ConstPetscScalarKokkosView3DType<DefaultMemorySpace>;
89: using ConstPetscScalarKokkosView4D = ConstPetscScalarKokkosView4DType<DefaultMemorySpace>;
91: using ConstPetscScalarKokkosViewHost = ConstPetscScalarKokkosViewType<HostMirrorMemorySpace>;
92: using ConstPetscScalarKokkosView1DHost = ConstPetscScalarKokkosView1DType<HostMirrorMemorySpace>;
93: using ConstPetscScalarKokkosView2DHost = ConstPetscScalarKokkosView2DType<HostMirrorMemorySpace>;
94: using ConstPetscScalarKokkosView3DHost = ConstPetscScalarKokkosView3DType<HostMirrorMemorySpace>;
95: using ConstPetscScalarKokkosView4DHost = ConstPetscScalarKokkosView4DType<HostMirrorMemorySpace>;
97: /* Shortcut types for OffsetViews in the default space and host space */
98: using PetscScalarKokkosOffsetView = PetscScalarKokkosOffsetViewType<DefaultMemorySpace>;
99: using PetscScalarKokkosOffsetView1D = PetscScalarKokkosOffsetView1DType<DefaultMemorySpace>;
100: using PetscScalarKokkosOffsetView2D = PetscScalarKokkosOffsetView2DType<DefaultMemorySpace>;
101: using PetscScalarKokkosOffsetView3D = PetscScalarKokkosOffsetView3DType<DefaultMemorySpace>;
102: using PetscScalarKokkosOffsetView4D = PetscScalarKokkosOffsetView4DType<DefaultMemorySpace>;
104: using PetscScalarKokkosOffsetViewHost = PetscScalarKokkosOffsetViewType<HostMirrorMemorySpace>;
105: using PetscScalarKokkosOffsetView1DHost = PetscScalarKokkosOffsetView1DType<HostMirrorMemorySpace>;
106: using PetscScalarKokkosOffsetView2DHost = PetscScalarKokkosOffsetView2DType<HostMirrorMemorySpace>;
107: using PetscScalarKokkosOffsetView3DHost = PetscScalarKokkosOffsetView3DType<HostMirrorMemorySpace>;
108: using PetscScalarKokkosOffsetView4DHost = PetscScalarKokkosOffsetView4DType<HostMirrorMemorySpace>;
110: using ConstPetscScalarKokkosOffsetView = ConstPetscScalarKokkosOffsetViewType<DefaultMemorySpace>;
111: using ConstPetscScalarKokkosOffsetView1D = ConstPetscScalarKokkosOffsetView1DType<DefaultMemorySpace>;
112: using ConstPetscScalarKokkosOffsetView2D = ConstPetscScalarKokkosOffsetView2DType<DefaultMemorySpace>;
113: using ConstPetscScalarKokkosOffsetView3D = ConstPetscScalarKokkosOffsetView3DType<DefaultMemorySpace>;
114: using ConstPetscScalarKokkosOffsetView4D = ConstPetscScalarKokkosOffsetView4DType<DefaultMemorySpace>;
116: using ConstPetscScalarKokkosOffsetViewHost = ConstPetscScalarKokkosOffsetViewType<HostMirrorMemorySpace>;
117: using ConstPetscScalarKokkosOffsetView1DHost = ConstPetscScalarKokkosOffsetView1DType<HostMirrorMemorySpace>;
118: using ConstPetscScalarKokkosOffsetView2DHost = ConstPetscScalarKokkosOffsetView2DType<HostMirrorMemorySpace>;
119: using ConstPetscScalarKokkosOffsetView3DHost = ConstPetscScalarKokkosOffsetView3DType<HostMirrorMemorySpace>;
120: using ConstPetscScalarKokkosOffsetView4DHost = ConstPetscScalarKokkosOffsetView4DType<HostMirrorMemorySpace>;
122: using PetscIntKokkosView = Kokkos::View<PetscInt *, DefaultMemorySpace>;
123: using PetscIntKokkosViewHost = Kokkos::View<PetscInt *, HostMirrorMemorySpace>;
124: using PetscIntKokkosDualView = Kokkos::DualView<PetscInt *>;
125: using PetscCountKokkosView = Kokkos::View<PetscCount *, DefaultMemorySpace>;
126: using PetscCountKokkosViewHost = Kokkos::View<PetscCount *, HostMirrorMemorySpace>;
128: // Sync a Kokkos::DualView to HostMirrorMemorySpace in execution space
129: // If is HostMirrorMemorySpace, fence the exec so that the data on host is immediately available.
130: template <typename Type>
131: static PetscErrorCode KokkosDualViewSyncHost(Kokkos::DualView<Type *> &v_dual, const Kokkos::DefaultExecutionSpace &exec)
132: {
133: size_t bytes = v_dual.extent(0) * sizeof(Type);
135: PetscFunctionBegin;
136: if (v_dual.need_sync_host()) {
137: PetscCallCXX(v_dual.sync_host(exec));
138: if (!std::is_same_v<DefaultMemorySpace, HostMirrorMemorySpace>) 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: PetscFunctionReturn(PETSC_SUCCESS);
145: }
147: template <typename Type>
148: static PetscErrorCode KokkosDualViewSyncDevice(Kokkos::DualView<Type *> &v_dual, const Kokkos::DefaultExecutionSpace &exec)
149: {
150: size_t bytes = v_dual.extent(0) * sizeof(Type);
152: PetscFunctionBegin;
153: if (v_dual.need_sync_device()) {
154: PetscCallCXX(v_dual.sync_device(exec));
155: if (!std::is_same_v<DefaultMemorySpace, HostMirrorMemorySpace>) PetscCall(PetscLogCpuToGpu(bytes));
156: }
157: PetscFunctionReturn(PETSC_SUCCESS);
158: }