AMSEOS Wrapper (#29)

* add AMSEOS wrapper * fix compile issues * minor * format + minor
LLNL · Nov 22, 2023 · 02f7475 · 02f7475
1 parent cd160ba
commit 02f7475
Show file tree

Hide file tree

Showing 7 changed files with 255 additions and 194 deletions.
diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -5,7 +5,7 @@
 
 file(GLOB_RECURSE AMS_CURRENT_EXAMPLE_INCLUDES "*.hpp")
 
-set(AMS_EXAMPLE_SRC ${MINIAPP_INCLUDES} main.cpp)
+set(AMS_EXAMPLE_SRC ${MINIAPP_INCLUDES} main.cpp app/eos_ams.cpp)
 
 
 function(ADDExec binary_name definitions)

diff --git a/examples/app/eos.hpp b/examples/app/eos.hpp
@@ -13,11 +13,21 @@
 template <typename FPType>
 class EOS
 {
-
 public:
-  virtual void Eval(const int length,
-                    const FPType **inputs,
-                    FPType **outputs) const = 0;
+#ifdef __ENABLE_PERFFLOWASPECT__
+  __attribute__((annotate("@critical_path(pointcut='around')")))
+#endif
+  virtual void
+  Eval(const int length, const FPType **inputs, FPType **outputs) const
+  {
+    Eval(length,
+         inputs[0],
+         inputs[1],
+         outputs[0],
+         outputs[1],
+         outputs[2],
+         outputs[3]);
+  }
 
   virtual void Eval(const int length,
                     const FPType *density,
@@ -36,20 +46,4 @@ class EOS
                                 FPType *bulkmod,
                                 FPType *temperature) const = 0;
 };
-
-template<typename FPType>
-void callBack(void *cls,
-              long elements,
-              const void *const *inputs,
-              void *const *outputs)
-{
-  static_cast<EOS<FPType> *>(cls)->Eval(elements,
-                                static_cast<const FPType *>(inputs[0]),
-                                static_cast<const FPType *>(inputs[1]),
-                                static_cast<FPType *>(outputs[0]),
-                                static_cast<FPType *>(outputs[1]),
-                                static_cast<FPType *>(outputs[2]),
-                                static_cast<FPType *>(outputs[3]));
-}
-
 #endif
diff --git a/examples/app/eos_ams.cpp b/examples/app/eos_ams.cpp
@@ -0,0 +1,97 @@
+/*
+ * Copyright 2021-2023 Lawrence Livermore National Security, LLC and other
+ * AMSLib Project Developers
+ *
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ */
+
+#include "eos_ams.hpp"
+
+#include <vector>
+
+template <typename FPType>
+void callBack(void *cls,
+              long elements,
+              const void *const *inputs,
+              void *const *outputs)
+{
+  static_cast<EOS<FPType> *>(cls)->Eval(elements,
+                                        static_cast<const FPType *>(inputs[0]),
+                                        static_cast<const FPType *>(inputs[1]),
+                                        static_cast<FPType *>(outputs[0]),
+                                        static_cast<FPType *>(outputs[1]),
+                                        static_cast<FPType *>(outputs[2]),
+                                        static_cast<FPType *>(outputs[3]));
+}
+
+
+template <typename FPType>
+AMSEOS<FPType>::AMSEOS(EOS<FPType> *model,
+                       const AMSDBType db_type,
+                       const AMSDType dtype,
+                       const AMSExecPolicy exec_policy,
+                       const AMSResourceType res_type,
+                       const AMSUQPolicy uq_policy,
+                       const int k_nearest,
+                       const int mpi_task,
+                       const int mpi_nproc,
+                       const double threshold,
+                       const char *surrogate_path,
+                       const char *uq_path,
+                       const char *db_path)
+    : model_(model)
+{
+  AMSConfig conf = {exec_policy,
+                    dtype,
+                    res_type,
+                    db_type,
+                    callBack<FPType>,
+                    (char *)surrogate_path,
+                    (char *)uq_path,
+                    (char *)db_path,
+                    threshold,
+                    uq_policy,
+                    k_nearest,
+                    mpi_task,
+                    mpi_nproc};
+
+  wf_ = AMSCreateExecutor(conf);
+}
+
+#ifdef __ENABLE_PERFFLOWASPECT__
+__attribute__((annotate("@critical_path(pointcut='around')")))
+#endif
+template <typename FPType>
+void AMSEOS<FPType>::Eval(const int length,
+                          const FPType *density,
+                          const FPType *energy,
+                          FPType *pressure,
+                          FPType *soundspeed2,
+                          FPType *bulkmod,
+                          FPType *temperature) const
+{
+  std::vector<const FPType *> inputs = {density, energy};
+  std::vector<FPType *> outputs = {pressure, soundspeed2, bulkmod, temperature};
+
+#ifdef __ENABLE_MPI__
+  AMSDistributedExecute(wf_,
+                        MPI_COMM_WORLD,
+                        (void *)model_,
+                        length,
+                        reinterpret_cast<const void **>(inputs.data()),
+                        reinterpret_cast<void **>(outputs.data()),
+                        inputs.size(),
+                        outputs.size());
+#else
+  AMSExecute(wf_,
+             (void *)model_,
+             length,
+             reinterpret_cast<const void **>(inputs.data()),
+             reinterpret_cast<void **>(outputs.data()),
+             inputs.size(),
+             outputs.size());
+#endif
+}
+
+template class AMSEOS<double>;
+template class AMSEOS<float>;
diff --git a/examples/app/eos_ams.hpp b/examples/app/eos_ams.hpp
@@ -0,0 +1,60 @@
+/*
+ * Copyright 2021-2023 Lawrence Livermore National Security, LLC and other
+ * AMSLib Project Developers
+ *
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ */
+
+#ifndef _AMS_EOS_HPP_
+#define _AMS_EOS_HPP_
+
+#include <stdexcept>
+
+#include "AMS.h"
+#include "eos.hpp"
+
+template <typename FPType>
+class AMSEOS : public EOS<FPType>
+{
+  AMSExecutor wf_;
+  EOS<FPType> *model_ = nullptr;
+
+public:
+  AMSEOS(EOS<FPType> *model,
+         const AMSDBType db_type,
+         const AMSDType dtype,
+         const AMSExecPolicy exec_policy,
+         const AMSResourceType res_type,
+         const AMSUQPolicy uq_policy,
+         const int k_nearest,
+         const int mpi_task,
+         const int mpi_nproc,
+         const double threshold,
+         const char *surrogate_path,
+         const char *uq_path,
+         const char *db_path);
+
+  virtual ~AMSEOS() { delete model_; }
+
+  void Eval(const int length,
+            const FPType *density,
+            const FPType *energy,
+            FPType *pressure,
+            FPType *soundspeed2,
+            FPType *bulkmod,
+            FPType *temperature) const override;
+
+  void Eval_with_filter(const int length,
+                        const FPType *density,
+                        const FPType *energy,
+                        const bool *filter,
+                        FPType *pressure,
+                        FPType *soundspeed2,
+                        FPType *bulkmod,
+                        FPType *temperature) const override
+  {
+    throw std::runtime_error("AMSEOS: Eval_with_filter is not implemented");
+  }
+};
+
+#endif  // _AMS_EOS_HPP_
diff --git a/examples/app/eos_constant_on_host.hpp b/examples/app/eos_constant_on_host.hpp
@@ -13,7 +13,7 @@
 
 #include "eos.hpp"
 
-template<typename FPType>
+template <typename FPType>
 class ConstantEOSOnHost : public EOS<FPType>
 {
   umpire::ResourceManager &rm_;
@@ -56,31 +56,16 @@ class ConstantEOSOnHost : public EOS<FPType>
   }
 
 #ifdef __ENABLE_PERFFLOWASPECT__
-    __attribute__((annotate("@critical_path(pointcut='around')")))
+  __attribute__((annotate("@critical_path(pointcut='around')")))
 #endif
-  void Eval(const int length,
-            const FPType **inputs,
-            FPType **outputs) const override
-  {
-    Eval(length,
-         inputs[0],
-         inputs[1],
-         outputs[0],
-         outputs[1],
-         outputs[2],
-         outputs[3]);
-  }
-
-#ifdef __ENABLE_PERFFLOWASPECT__
-    __attribute__((annotate("@critical_path(pointcut='around')")))
-#endif
-  void Eval(const int length,
-            const FPType *density,
-            const FPType *energy,
-            FPType *pressure,
-            FPType *soundspeed2,
-            FPType *bulkmod,
-            FPType *temperature) const override
+  void
+  Eval(const int length,
+       const FPType *density,
+       const FPType *energy,
+       FPType *pressure,
+       FPType *soundspeed2,
+       FPType *bulkmod,
+       FPType *temperature) const override
   {
     Eval_with_filter(length,
                      density,
@@ -93,16 +78,17 @@ class ConstantEOSOnHost : public EOS<FPType>
   }
 
 #ifdef __ENABLE_PERFFLOWASPECT__
-    __attribute__((annotate("@critical_path(pointcut='around')")))
+  __attribute__((annotate("@critical_path(pointcut='around')")))
 #endif
-  void Eval_with_filter(const int length,
-                        const FPType *density,
-                        const FPType *energy,
-                        const bool *filter,
-                        FPType *pressure,
-                        FPType *soundspeed2,
-                        FPType *bulkmod,
-                        FPType *temperature) const override
+  void
+  Eval_with_filter(const int length,
+                   const FPType *density,
+                   const FPType *energy,
+                   const bool *filter,
+                   FPType *pressure,
+                   FPType *soundspeed2,
+                   FPType *bulkmod,
+                   FPType *temperature) const override
   {
     auto plt = GetPlatform((void *)density);
     const char *res_name = GetResourceName(plt);

diff --git a/examples/app/eos_idealgas.hpp b/examples/app/eos_idealgas.hpp
@@ -13,7 +13,7 @@
 
 //! Ideal Gas EOS
 //! Code given by Thomas Stitt
-template<typename FPType>
+template <typename FPType>
 class IdealGas : public EOS<FPType>
 {
   const FPType gamma_;
@@ -26,15 +26,16 @@ class IdealGas : public EOS<FPType>
   }
 
 #ifdef __ENABLE_PERFFLOWASPECT__
-   __attribute__((annotate("@critical_path(pointcut='around')")))
+  __attribute__((annotate("@critical_path(pointcut='around')")))
 #endif
-  void Eval(const int length,
-            const FPType *density,
-            const FPType *energy,
-            FPType *pressure,
-            FPType *soundspeed2,
-            FPType *bulkmod,
-            FPType *temperature) const override
+  void
+  Eval(const int length,
+       const FPType *density,
+       const FPType *energy,
+       FPType *pressure,
+       FPType *soundspeed2,
+       FPType *bulkmod,
+       FPType *temperature) const override
   {
     const FPType gamma = gamma_;
     const FPType specific_heat = specific_heat_;
@@ -48,14 +49,18 @@ class IdealGas : public EOS<FPType>
     });
   }
 
-  void Eval_with_filter(const int length,
-                        const FPType *density,
-                        const FPType *energy,
-                        const bool *filter,
-                        FPType *pressure,
-                        FPType *soundspeed2,
-                        FPType *bulkmod,
-                        FPType *temperature) const override
+#ifdef __ENABLE_PERFFLOWASPECT__
+  __attribute__((annotate("@critical_path(pointcut='around')")))
+#endif
+  void
+  Eval_with_filter(const int length,
+                   const FPType *density,
+                   const FPType *energy,
+                   const bool *filter,
+                   FPType *pressure,
+                   FPType *soundspeed2,
+                   FPType *bulkmod,
+                   FPType *temperature) const override
   {
     const FPType gamma = gamma_;
     const FPType specific_heat = specific_heat_;
@@ -70,22 +75,5 @@ class IdealGas : public EOS<FPType>
       }
     });
   }
-
-#ifdef __ENABLE_PERFFLOWASPECT__
-   __attribute__((annotate("@critical_path(pointcut='around')")))
-#endif
-  void Eval(const int length,
-            const FPType **inputs,
-            FPType **outputs) const override
-  {
-    Eval(length,
-         inputs[0],
-         inputs[1],
-         outputs[0],
-         outputs[1],
-         outputs[2],
-         outputs[3]);
-  }
 };
-
 #endif