1 // Copyright (c) 2010-2021, Lawrence Livermore National Security, LLC. Produced 2 // at the Lawrence Livermore National Laboratory. All Rights reserved. See files 3 // LICENSE and NOTICE for details. LLNL-CODE-806117. 4 // 5 // This file is part of the MFEM library. For more information and source code 6 // availability visit https://mfem.org. 7 // 8 // MFEM is free software; you can redistribute it and/or modify it under the 9 // terms of the BSD-3 license. We welcome feedback and contributions, see file 10 // CONTRIBUTING.md for details. 11 12 #ifndef MFEM_TMOP_PA_HPP 13 #define MFEM_TMOP_PA_HPP 14 15 #include "../../config/config.hpp" 16 #include "../../linalg/dtensor.hpp" 17 18 #include "../kernels.hpp" 19 20 #include <unordered_map> 21 22 namespace mfem 23 { 24 25 namespace kernels 26 { 27 28 template <typename K> class KernelMap; 29 30 /// Instances 31 template<class K, int N = K::N> struct Instances 32 { Fillmfem::kernels::Instances33 static void Fill(KernelMap<K> &map) 34 { 35 map.template Emplace<N-1>(); 36 Instances<K,N-1>::Fill(map); 37 } 38 }; 39 40 // terminal case 41 template<class K> struct Instances<K,1> 42 { Fillmfem::kernels::Instances43 static void Fill(KernelMap<K> &map) { map.template Emplace<0>(); } 44 }; 45 46 /// KernelMap class which creates an unordered_map of the Keys/Kernels 47 template<class K> 48 class KernelMap 49 { 50 using Key_t = typename K::Key_t; 51 using Return_t = typename K::Return_t; 52 using Kernel_t = typename K::Kernel_t; 53 using map_t = std::unordered_map<Key_t, Kernel_t>; 54 map_t map; 55 56 public: 57 // Fill all the map with the Keys/Kernels KernelMap()58 KernelMap() { Instances<K>::Fill(*this); } 59 Find(const Key_t id)60 bool Find(const Key_t id) { return map.find(id) != map.end(); } 61 At(const Key_t id)62 Kernel_t At(const Key_t id) { return map.at(id); } 63 Emplace()64 template<int N> void Emplace() 65 { 66 constexpr Key_t key = K::template GetKey<N>(); 67 constexpr Kernel_t ker = K::template GetKer<key>(); 68 map.emplace(key, ker); 69 } 70 }; 71 72 // ///////////////////////////////////////////////////////////////////////////// 73 // MFEM_REGISTER_TMOP_KERNELS macro: 74 // - the first argument (return_t) is the return type of the kernel 75 // - the second argument (kernel) is the name of the kernel 76 // - the arguments of the kernel (...) captured as __VA_ARGS__ 77 // 78 // This call will output the followings: 79 // 1. forward declaration of the kernel 80 // 2. kernel pointer declaration 81 // 3. struct K##name##_T definition which holds the keys/kernels map 82 // 4. KernelMap definition of the current kernel 83 // 5. the kernel signature by re-using all the arguments 84 // 85 // ///////////////////////////////////////////////////////////////////////////// 86 // For example: 87 // MFEM_REGISTER_TMOP_KERNELS(void, Name, 88 // const int NE, 89 // const Array<double> &b, 90 // Vector &diagonal, 91 // const int d1d, 92 // const int q1d) {...} 93 // 94 // The resulting code would be: 95 // 96 // 1. forward declaration of the kernel 97 // template<int T_D1D = 0, int T_Q1D = 0, int T_MAX = 0> 98 // void Name(const int NE, 99 // const Array<double> &b, 100 // Vector &diagonal, 101 // const int d1d, 102 // const int q1d); 103 // 104 // 2. kernel pointer declaration 105 // typedef void (*Name_p)(const int NE, 106 // const Array<double> &b, 107 // Vector &diagonal, 108 // const int d1d, 109 // const int q1d); 110 // 111 // 3. struct K##Name##_T definition which holds the keys/kernels instance 112 // struct KName_T 113 // { 114 // static const int N = 14; 115 // using Key_t = int; 116 // using Return_t = void; 117 // using Kernel_t = Name_p; 118 // template<Key_t I> static constexpr Key_t GetKey() noexcept 119 // { 120 // return I==0 ? 0x22 : I==1 ? 0x23 : I==2 ? 0x24 : I==3 ? 0x25 : 121 // I==4 ? 0x26 : I== 5 ? 0x33 : I==6 ? 0x34 : I==7 ? 0x35 : 122 // I==8 ? 0x36 : I==9 ? 0x44 : I==10 ? 0x45 : I==11 ? 0x46 : 123 // I==12 ? 0x55 : I==13 ? 0x56 : 0; 124 // } 125 // template<Key_t K> static constexpr Kernel_t GetKer() noexcept 126 // { 127 // return &AssembleDiagonalPA_Kernel_2D<(K>>4)&0xF, K&0xF>; 128 // } 129 // }; 130 // 131 // 4. KernelMap definition of the current kernel 132 // static kernels::KernelMap<KName_T> KName; 133 // 134 // 5. the kernel signature by re-using all the arguments 135 // template<int T_D1D, int T_Q1D, int T_MAX> 136 // void Name(const int NE, 137 // const Array<double> &b, 138 // Vector &diagonal, 139 // const int d1d, 140 // const int q1d) {...} 141 142 // ///////////////////////////////////////////////////////////////////////////// 143 // All of which allows to launch the kernel with a specific id ((D1D<<4)|Q1D). 144 // 145 // For example, a MFEM_LAUNCH_TMOP_KERNEL(Name,id,NE,B,D); call would result in: 146 // 147 // if (KName.Find(id)) { return KName.At(id)(NE,B,D,0,0); } 148 // else 149 // { 150 // constexpr int T_MAX = 4; 151 // const int D1D = (id>>4)&0xF, Q1D = id&0xF; 152 // MFEM_VERIFY(D1D <= MAX_D1D && Q1D <= MAX_Q1D, "Max size error!"); 153 // return Name<0,0,T_MAX>(NE,B,D,D1D,Q1D); 154 // }; 155 156 #define MFEM_REGISTER_TMOP_KERNELS(return_t, kernel, ...) \ 157 template<int T_D1D = 0, int T_Q1D = 0, int T_MAX = 0> \ 158 return_t kernel(__VA_ARGS__);\ 159 typedef return_t (*kernel##_p)(__VA_ARGS__);\ 160 struct K##kernel##_T {\ 161 static const int N = 14;\ 162 using Key_t = int;\ 163 using Return_t = return_t;\ 164 using Kernel_t = kernel##_p;\ 165 template<Key_t I> static constexpr Key_t GetKey() noexcept { return \ 166 I==0 ? 0x22 : I==1 ? 0x23 : I==2 ? 0x24 : I==3 ? 0x25 : I==4 ? 0x26 :\ 167 I==5 ? 0x33 : I==6 ? 0x34 : I==7 ? 0x35 : I==8 ? 0x36 :\ 168 I==9 ? 0x44 : I==10 ? 0x45 : I==11 ? 0x46 :\ 169 I==12 ? 0x55 : I==13 ? 0x56 : 0; }\ 170 template<Key_t K> static constexpr Kernel_t GetKer() noexcept\ 171 { return &kernel<(K>>4)&0xF, K&0xF>; }\ 172 };\ 173 static kernels::KernelMap<K##kernel##_T> K##kernel;\ 174 template<int T_D1D, int T_Q1D, int T_MAX> return_t kernel(__VA_ARGS__) 175 176 // MFEM_LAUNCH_TMOP_KERNEL macro 177 // This macro will try to find and launch the kernel with the id key and 178 // the templated arguments. 179 // If not, it will fall back to the kernel with the standard arguments. 180 #define MFEM_LAUNCH_TMOP_KERNEL(kernel, id, ...)\ 181 if (K##kernel.Find(id)) { return K##kernel.At(id)(__VA_ARGS__,0,0); }\ 182 else {\ 183 constexpr int T_MAX = 4;\ 184 const int D1D = (id>>4)&0xF, Q1D = id&0xF;\ 185 MFEM_VERIFY(D1D <= MAX_D1D && Q1D <= MAX_Q1D, "Max size error!");\ 186 return kernel<0,0,T_MAX>(__VA_ARGS__,D1D,Q1D); } 187 188 } // namespace kernels 189 190 } // namespace mfem 191 192 #endif // MFEM_TMOP_PA_HPP 193