Adds -rp parameter and option to run ML model on GPU and physics module on CPU

examples/app/eos_idealgas.hpp

@@ -18,10 +18,11 @@ class IdealGas : public EOS<FPType>
 {
   const FPType gamma_;
   const FPType specific_heat_;
+  const int repeat;
 
 public:
-  IdealGas(FPType gamma, FPType specific_heat)
-      : gamma_(gamma), specific_heat_(specific_heat)
+  IdealGas(FPType gamma, FPType specific_heat, int repeat)
+      : gamma_(gamma), specific_heat_(specific_heat), repeat(repeat)
   {
   }
 
@@ -39,14 +40,15 @@ class IdealGas : public EOS<FPType>
   {
     const FPType gamma = gamma_;
     const FPType specific_heat = specific_heat_;
-
-    using mfem::ForallWrap;
-    MFEM_FORALL(i, length, {
-      pressure[i] = (gamma - 1) * density[i] * energy[i];
-      soundspeed2[i] = gamma * (gamma - 1) * energy[i];
-      bulkmod[i] = gamma * pressure[i];
-      temperature[i] = energy[i] / specific_heat;
+    for ( int r = 0; r < repeat; r++){
+      using mfem::ForallWrap;
+      MFEM_FORALL(i, length, {
+        pressure[i] = (gamma - 1) * density[i] * energy[i];
+        soundspeed2[i] = gamma * (gamma - 1) * energy[i];
+        bulkmod[i] = gamma * pressure[i];
+        temperature[i] = energy[i] / specific_heat;
     });
+    }
   }
 
 #ifdef __ENABLE_PERFFLOWASPECT__

examples/main.cpp

@@ -166,7 +166,8 @@ int run(const char *device_name,
         const char *model_path,
         const char *db_config,
         bool lbalance,
-        int k_nearest)
+        int k_nearest,
+        int repeats)
 {
   // -------------------------------------------------------------------------
   // setup
@@ -175,7 +176,10 @@ int run(const char *device_name,
   CALIPER(mgr.start();)
   CALIPER(CALI_MARK_BEGIN("Setup");)
 
-  const bool use_device = std::strcmp(device_name, "cpu") != 0;
+  const bool physics_use_device = ((std::strcmp(device_name, "gpu_cpu") == 0) || std::strcmp(device_name, "gpu_gpu") == 0);
+  const bool ml_use_device = ((std::strcmp(device_name, "gpu_cpu") == 0) || std::strcmp(device_name, "gpu_gpu") == 0);
+
+
   AMSDBType dbType = AMSDBType::None;
   if (std::strcmp(db_type, "csv") == 0) {
     dbType = AMSDBType::CSV;
@@ -250,7 +254,7 @@ int run(const char *device_name,
 
 
   mfem::MemoryManager::SetUmpireHostAllocatorName(physics_host_alloc.c_str());
-  if (use_device) {
+  if (physics_use_device) {
     mfem::MemoryManager::SetUmpireDeviceAllocatorName(
         physics_device_alloc.c_str());
   }
@@ -261,7 +265,7 @@ int run(const char *device_name,
   if (strcmp(pool, "default") != 0) {
     AMSSetAllocator(AMSResourceType::HOST, ams_host_alloc.c_str());
 
-    if (use_device) {
+    if ( physics_use_device || ml_use_device ) {
       AMSSetAllocator(AMSResourceType::DEVICE, ams_device_alloc.c_str());
       AMSSetAllocator(AMSResourceType::PINNED, ams_pinned_alloc.c_str());
     }
@@ -270,7 +274,13 @@ int run(const char *device_name,
   mfem::Device::SetMemoryTypes(mfem::MemoryType::HOST_UMPIRE,
                                mfem::MemoryType::DEVICE_UMPIRE);
 
-  mfem::Device device(device_name);
+  mfem::Device device;
+
+  if ( physics_use_device )
+    device.Configure("cuda");
+  else
+    device.Configure("cpu");
+
   std::cout << std::endl;
   device.Print();
   std::cout << std::endl;
@@ -332,8 +342,12 @@ int run(const char *device_name,
 
   db_path = (strlen(db_config) > 0) ? db_config : nullptr;
 
-  AMSResourceType ams_device = AMSResourceType::HOST;
-  if (use_device) ams_device = AMSResourceType::DEVICE;
+  AMSResourceType ams_physics_device = AMSResourceType::HOST;
+  if (physics_use_device) ams_physics_device = AMSResourceType::DEVICE;
+
+  AMSResourceType ams_ml_device = AMSResourceType::HOST;
+  if (ml_use_device) ams_ml_device = AMSResourceType::DEVICE;
+
   AMSExecPolicy ams_loadBalance = AMSExecPolicy::UBALANCED;
   if (lbalance) ams_loadBalance = AMSExecPolicy::BALANCED;
 #else
@@ -347,7 +361,7 @@ int run(const char *device_name,
   for (int mat_idx = 0; mat_idx < num_mats; ++mat_idx) {
     EOS<TypeValue> *base;
     if (eos_name == std::string("ideal_gas")) {
-      base = new IdealGas<TypeValue>(1.6, 1.4);
+      base = new IdealGas<TypeValue>(1.6, 1.4, repeats);
     } else if (eos_name == std::string("constant_host")) {
       base = new ConstantEOSOnHost<TypeValue>(physics_host_alloc.c_str(), 1.0);
     } else {
@@ -605,7 +619,7 @@ int main(int argc, char **argv)
     std::cout.setstate(std::ios::failbit);
   }
 
-  const char *device_name = "cpu";
+  const char *device_name = "cpu_cpu";
   const char *eos_name = "ideal_gas";
   const char *model_path = "";
   const char *hdcache_path = "";
@@ -634,6 +648,7 @@ int main(int argc, char **argv)
   double avg = 0.5;
   double stdDev = 0.2;
   bool reqDB = false;
+  int repeats = 1;
   const char *pool = "default";
 
 #ifdef __ENABLE_DB__
@@ -646,7 +661,7 @@ int main(int argc, char **argv)
   // setup command line parser
   // -------------------------------------------------------------------------
   mfem::OptionsParser args(argc, argv);
-  args.AddOption(&device_name, "-d", "--device", "Device config string");
+  args.AddOption(&device_name, "-d", "--device", "device(physics)_device(model), for example: cpu_gpu means execute physics on the cpu and the model on gpu");
 
   // set precision
   args.AddOption(&precision_opt,
@@ -752,6 +767,8 @@ int main(int argc, char **argv)
   args.AddOption(
       &verbose, "-v", "--verbose", "-qu", "--quiet", "Print extra stuff");
 
+  args.AddOption(&repeats, "-rp", "--repeats", "Number of repeats to perform on physics eval");
+
   args.AddOption(&pool,
                  "-ptype",
                  "--pool-type",
@@ -844,7 +861,8 @@ int main(int argc, char **argv)
                      model_path,
                      db_config,
                      lbalance,
-                     k_nearest);
+                     k_nearest,
+                     repeats);
   else if (precision == AMSDType::Double)
     ret = run<double>(device_name,
                       db_type,
@@ -870,7 +888,8 @@ int main(int argc, char **argv)
                       model_path,
                       db_config,
                       lbalance,
-                      k_nearest);
+                      k_nearest,
+                      repeats);
   else {
     std::cerr << "Invalid precision " << precision_opt << "\n";
     return -1;

src/AMSlib/AMS.cpp

@@ -87,7 +87,8 @@ AMSExecutor AMSCreateExecutor(const AMSConfig config)
                                      config.SPath,
                                      config.DBPath,
                                      config.dbType,
-                                     config.device,
+                                     config.physics_device,
+                                     config.ml_device,
                                      config.threshold,
                                      config.uqPolicy,
                                      config.nClusters,
@@ -105,7 +106,8 @@ AMSExecutor AMSCreateExecutor(const AMSConfig config)
                                     config.SPath,
                                     config.DBPath,
                                     config.dbType,
-                                    config.device,
+                                    config.physics_device,
+                                    config.ml_device,
                                     static_cast<float>(config.threshold),
                                     config.uqPolicy,
                                     config.nClusters,

src/AMSlib/include/AMS.h

@@ -71,7 +71,8 @@ enum struct AMSUQPolicy {
 typedef struct ams_conf {
   const AMSExecPolicy ePolicy;
   const AMSDType dType;
-  const AMSResourceType device;
+  const AMSResourceType physics_device;
+  const AMSResourceType ml_device;
   const AMSDBType dbType;
   AMSPhysicFn cBack;
   char *SPath;

src/AMSlib/ml/surrogate.hpp

@@ -156,7 +156,7 @@ class SurrogateModel
   inline void _load(const std::string& model_path,
                     const std::string& device_name)
   {
-    DBG(Surrogate, "Using model at double precision");
+    DBG(Surrogate, "Using model(%s) at double precision at device %s", model_path.c_str(), device_name.c_str());
     _load_torch(model_path, torch::Device(device_name), torch::kFloat64);
   }
 
@@ -166,7 +166,7 @@ class SurrogateModel
   inline void _load(const std::string& model_path,
                     const std::string& device_name)
   {
-    DBG(Surrogate, "Using model at single precision");
+    DBG(Surrogate, "Using model (%s) at single precision at device %s", model_path.c_str(), device_name.c_str())
     _load_torch(model_path, torch::Device(device_name), torch::kFloat32);
   }