Update BO to new API

LLNL · Jan 3, 2025 · 15efb25 · 15efb25
1 parent 2104982
commit 15efb25
Show file tree

Hide file tree

Showing 4 changed files with 94 additions and 40 deletions.
diff --git a/examples/bnm_opt/CMakeLists.txt b/examples/bnm_opt/CMakeLists.txt
@@ -3,25 +3,65 @@
 #
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-set(binomial_options_src binomial_options.cpp kernel.cpp)
-
-function(ADDExec binary_name definitions)
-  target_include_directories(${binary_name} PUBLIC "${AMS_EXAMPLE_INCLUDES}")
-  target_compile_definitions(${binary_name} PRIVATE ${definitions})
-  target_compile_definitions(${binary_name} PRIVATE ${AMS_EXAMPLE_DEFINES})
-  target_link_directories(${binary_name} PRIVATE ${AMS_EXAMPLE_LIB_DIRS})
-  target_link_libraries(${binary_name} PUBLIC AMS ${AMS_EXAMPLE_LIBRARIES})
-
-  if (WITH_CUDA)
-      set_source_files_properties(kernel.cpp PROPERTIES LANGUAGE CUDA)
-      set_target_properties(${binary_name} PROPERTIES CUDA_ARCHITECTURES "${AMS_CUDA_ARCH}")
+
+cmake_minimum_required(VERSION 3.10)
+cmake_policy(SET CMP0104 NEW)
+
+# Define the project
+project(BOptions LANGUAGES CXX)
+
+set(AMS_EXAMPLE_SRC binomial_options.cpp kernel.cpp)
+
+
+option(WITH_MPI            "Option to enable MPI" OFF)
+option(WITH_CALIPER        "Use Caliper for Profiling" OFF)
+
+set(CMAKE_CXX_STANDARD 14)      # Enable C++14
+set(CMAKE_CXX_STANDARD_REQUIRED ON) # Require the specified standard
+
+enable_language(CUDA)
+find_package(CUDAToolkit REQUIRED)
+
+if (WITH_CALIPER)
+  find_package(caliper REQUIRED)
+endif()
+
+if (WITH_MPI)
+  find_package(MPI REQUIRED)
+endif()
+
+function(ADDExec binary_name use_ams)
+  target_link_libraries(${binary_name} PRIVATE ${MFEM_LIBRARIES})
+  target_include_directories(${binary_name} PRIVATE ${MFEM_INCLUDE_DIRS})
+  # we always use ams to avoid if def conflicts.
+  # but we only enable it for the use_ams case
+  find_package(AMS REQUIRED)
+  target_link_libraries(${binary_name} PUBLIC AMS::AMS)
+  if (${use_ams})
+    target_compile_definitions(${binary_name} PRIVATE USE_AMS)
+  endif()
+
+  if (WITH_MPI)
+    target_link_libraries(${binary_name} PRIVATE MPI::MPI_CXX)
+    target_compile_definitions(${binary_name} PRIVATE "__ENABLE_MPI__")
   endif()
+
+
+  target_link_libraries(${binary_name} PRIVATE 
+    CUDA::cudart
+    CUDA::cublas
+    CUDA::cusparse
+    CUDA::curand 
+  )
+  set_source_files_properties(kernel.cpp PROPERTIES LANGUAGE CUDA)
+  set_target_properties(${binary_name} PROPERTIES CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}")
+  target_compile_definitions(${binary_name} PRIVATE "__ENABLE_CUDA__")
 endfunction()
 
-add_executable(no_ams_bo ${binomial_options_src} ${MINIAPP_INCLUDES})
-ADDExec(no_ams_bo minibude "${AMS_EXAMPLE_DEFINES}")
+add_executable(no_ams_boptions ${AMS_EXAMPLE_SRC} ${MINIAPP_INCLUDES})
+ADDExec(no_ams_boptions FALSE)
+
+add_executable(ams_boptions ${AMS_EXAMPLE_SRC} ${MINIAPP_INCLUDES})
+ADDExec(ams_boptions TRUE)
 
-list(APPEND AMS_EXAMPLE_DEFINES "-DUSE_AMS")
-add_executable(ams_bo ${binomial_options_src} ${MINIAPP_INCLUDES})
-ADDExec(ams_bo "${AMS_EXAMPLE_DEFINES}")
 
diff --git a/examples/bnm_opt/binomial_options.cpp b/examples/bnm_opt/binomial_options.cpp
@@ -253,7 +253,7 @@ int main(int argc, char **argv)
 #endif
 
   if (argc != 5) {
-    std::cout << "USAGE: " << argv[0] << " num-options batch_size";
+    std::cout << "USAGE: " << argv[0] << " num-options batch_size num_fractions fraction_id";
     return EXIT_FAILURE;
   }
 

diff --git a/examples/bnm_opt/kernel.cpp b/examples/bnm_opt/kernel.cpp
@@ -26,6 +26,8 @@
 #include <AMS.h>
 #endif
 
+using namespace ams;
+
 
 // Overloaded shortcut functions for different precision modes
 #ifndef DOUBLE_PRECISION
@@ -169,16 +171,14 @@ BinomialOptions::BinomialOptions(unsigned int batchSize,
 
 #ifdef USE_AMS
   const char *model_name = std::getenv("BO_MODEL_NAME");
+  std::cout << "Model name is " << model_name << "\n";
   if (model_name) {
     model = AMSQueryModel(model_name);
   } else {
     model = AMSQueryModel("binomialOptions");
   }
 
   wf = AMSCreateExecutor(model,
-                         AMSDType::AMS_DOUBLE,
-                         AMSResourceType::AMS_DEVICE,
-                         (AMSPhysicFn)(BinomialOptions::AMSRun),
                          rank,
                          worldSize);
 #endif
@@ -219,18 +219,38 @@ void BinomialOptions::run(real *callValue,
   cudaMemcpy(d_X, _X, sizeof(real) * optN, cudaMemcpyHostToDevice);
 
 #ifdef USE_AMS
-  std::vector<const real *> inputs({(const real *)d_S,
-                                    (const real *)d_X,
-                                    (const real *)d_R,
-                                    (const real *)d_V,
-                                    (const real *)d_T});
+
+  SmallVector<AMSTensor> inputs;
+  SmallVector<AMSTensor> inout;
+  SmallVector<AMSTensor> outputs;
+  inputs.push_back(std::move(AMSTensor::view(d_S, {static_cast<long>(optN), 1L}, {1, 1}, AMSResourceType::AMS_DEVICE)));
+  inputs.push_back(std::move(AMSTensor::view(d_X, {static_cast<long>(optN), 1L}, {1, 1}, AMSResourceType::AMS_DEVICE)));
+  inputs.push_back(std::move(AMSTensor::view(d_R, {static_cast<long>(optN), 1L}, {1, 1}, AMSResourceType::AMS_DEVICE)));
+  inputs.push_back(std::move(AMSTensor::view(d_V, {static_cast<long>(optN), 1L}, {1, 1}, AMSResourceType::AMS_DEVICE)));
+  inputs.push_back(std::move(AMSTensor::view(d_T, {static_cast<long>(optN), 1L}, {1, 1}, AMSResourceType::AMS_DEVICE)));
+
+
+  outputs.push_back(std::move(AMSTensor::view(d_CallValue, {static_cast<long>(optN), 1}, {1, 1}, AMSResourceType::AMS_DEVICE)));
+
+  EOSLambda OrigComputation = [&, this](const SmallVector<AMSTensor> &ams_ins,
+                                        SmallVector<AMSTensor> &ams_inouts,
+                                        SmallVector<AMSTensor> &ams_outs) {
+  binomialOptionsGPU(ams_outs[0].data<real>(),
+                     ams_ins[0].data<real>(),
+                     ams_ins[1].data<real>(),
+                     ams_ins[2].data<real>(),
+                     ams_ins[3].data<real>(),
+                     ams_ins[4].data<real>(),
+                     d_vDt,
+                     d_puByDf,
+                     d_pdByDf,
+                     ams_outs[0].shape()[0]);
+  };
+
+
   AMSExecute(wf,
-             (void *)this,
-             optN,
-             reinterpret_cast<const void **>(inputs.data()),
-             reinterpret_cast<void **>(&d_CallValue),
-             inputs.size(),
-             1);
+             OrigComputation,
+             inputs, inout, outputs);
 #else
   binomialOptionsGPU(
       d_CallValue, d_S, d_X, d_R, d_V, d_T, d_puByDf, d_pdByDf, d_vDt, optN);

diff --git a/examples/bnm_opt/kernel.hpp b/examples/bnm_opt/kernel.hpp
@@ -2,9 +2,7 @@
 
 #include "realtype.h"
 
-#ifdef USE_AMS
 #include <AMS.h>
-#endif
 
 class BinomialOptions
 {
@@ -18,10 +16,8 @@ class BinomialOptions
   real *d_X;
   real *d_CallValue;
 
-#ifdef USE_AMS
-  AMSCAbstrModel model;
-  AMSExecutor wf;
-#endif
+  ams::AMSCAbstrModel model;
+  ams::AMSExecutor wf;
 
 public:
   BinomialOptions(unsigned int batchSize, int rank, int worldSize);
@@ -33,9 +29,7 @@ class BinomialOptions
            real *_T,
            size_t optN);
 
-#ifdef USE_AMS
   static void AMSRun(void *cls, long numOptions, void **inputs, void **outputs);
-#endif
 
   ~BinomialOptions();
 };