Isolate AMS from PyTorch

.github/workflows/ci.yml

@@ -18,12 +18,13 @@ jobs:
       matrix:
         with_caliper: [ON, OFF]
         with_mpi: [ON, OFF]
+        with_torch: [ON, OFF]
 
     steps:
       - uses: actions/checkout@v7
       - name: Ensure git safe directory
         run: git config --global --add safe.directory '*'
-      - name: Build ENABLE_CALIPER=${{ matrix.with_caliper }} ENABLE_MPI=${{ matrix.with_mpi }}
+      - name: Build ENABLE_TORCH=${{ matrix.with_torch }} ENABLE_CALIPER=${{ matrix.with_caliper }} ENABLE_MPI=${{ matrix.with_mpi }}
         shell: bash -l {0}
         run: |
           source /etc/profile
@@ -42,6 +43,7 @@ jobs:
             -DENABLE_TESTS=On \
             -DAMS_ENABLE_DEBUG=On \
             -DENABLE_WORKFLOW=Off \
+            -DENABLE_TORCH=${{ matrix.with_torch }} \
             -DTorch_DIR=$AMS_TORCH_PATH \
             -Dcaliper_DIR=$AMS_CALIPER_PATH \
             -DAMS_FMT_DIR=$AMS_FMT_DIR \

.gitlab/jobs/dane.yml

@@ -61,8 +61,9 @@ build-run-dane:
 .build-variants:
   parallel:
     matrix:
-      - WITH_MPI: ["on", "off"]
-        WITH_WORKFLOW: ["on", "off"]
+      - WITH_MPI: ["ON", "OFF"]
+        WITH_WORKFLOW: ["ON", "OFF"]
+        WITH_TORCH: ["ON", "OFF"]
 
 build-run-dane:
   extends: [.base-job, .build-variants]

.gitlab/jobs/tioga.yml

@@ -36,8 +36,9 @@ variables:
 .build-variants:
   parallel:
     matrix:
-      - WITH_MPI: ["on", "off"]
-        WITH_WORKFLOW: ["on", "off"]
+      - WITH_MPI: ["ON", "OFF"]
+        WITH_WORKFLOW: ["ON", "OFF"]
+        WITH_TORCH: ["ON", "OFF"]
 
 build-run-tioga:
   extends: [.base-job, .build-variants]

.gitlab/jobs/tuolumne.yml

@@ -40,8 +40,9 @@ build-run-tuolumne:
 .build-variants:
   parallel:
     matrix:
-      - WITH_MPI: ["on", "off"]
-        WITH_WORKFLOW: ["on", "off"]
+      - WITH_MPI: ["ON", "OFF"]
+        WITH_WORKFLOW: ["ON", "OFF"]
+        WITH_TORCH: ["ON", "OFF"]
 
 build-run-tuolumne:
   extends: [.base-job, .build-variants]

CMakeLists.txt

@@ -49,6 +49,7 @@ endif()
 option(ENABLE_WORKFLOW "Install python drivers used by the outer workflow" OFF)
 option(ENABLE_RMQ "Use RabbitMQ as a database back end" OFF)
 option(ENABLE_PERFFLOWASPECT "Use PerfFlowAspect for profiling" OFF)
+option(ENABLE_TORCH "Enable PyTorch ML inference support" ON)
 option(AMS_ENABLE_DEBUG "Enable verbose AMS messages" OFF)
 option(AMS_INSTALL_FLUX_PYTHON
        "Install AMS Workflow Python package with the flux-python optional dependency"
@@ -240,6 +241,7 @@ if (AMS_DEFER_STATIC_TPL_RESOLUTION AND NOT BUILD_SHARED_LIBS)
           "AMS_DEFER_STATIC_TPL_RESOLUTION is ignored when BUILD_SHARED_LIBS=OFF.")
 endif()
 
+# set(THREADS_PREFER_PTHREAD_FLAG ON)
 find_package(Threads REQUIRED)
 
 include(cmake/FetchAndAddFmt.cmake)
@@ -289,7 +291,17 @@ if (ENABLE_CUDA)
   find_package(CUDAToolkit REQUIRED)
   list(APPEND AMS_APP_DEFINES "__AMS_ENABLE_CUDA__")
 elseif (ENABLE_HIP)
-  find_package(HIP REQUIRED)
+  find_package(hip REQUIRED)
+  find_package(hiprtc QUIET) # can be needed down the line
+
+  # a bit ugly here
+  if (NOT TARGET hiprtc::hiprtc)
+    message(STATUS "hiprtc not found, linking directly to ${ROCM_PATH}/lib/libhiprtc.so")
+    add_library(hiprtc::hiprtc UNKNOWN IMPORTED)
+    set_target_properties(hiprtc::hiprtc PROPERTIES
+      IMPORTED_LOCATION "${ROCM_PATH}/lib/libhiprtc.so")
+  endif()
+
   list(APPEND AMS_APP_DEFINES "__AMS_ENABLE_HIP__")
   if (DEFINED ROCM_PATH)
     string(APPEND CMAKE_CXX_FLAGS "-I${ROCM_PATH}/include/")
@@ -390,11 +402,46 @@ if (ENABLE_RMQ)
   find_package(libevent REQUIRED)
 endif()
 
-find_package(Torch REQUIRED)
-# This is annoying, torch populates all my cuda flags
-# and resets them
-set(CMAKE_CUDA_FLAGS "")
-set(CMAKE_CUDA_ARCHITECTURES ON)
+# ------------------------------------------------------------------------------
+if (ENABLE_TORCH)
+  find_package(Torch REQUIRED)
+  # This is annoying, torch populates all my cuda flags
+  # and resets them
+  set(CMAKE_CUDA_FLAGS "")
+  set(CMAKE_CUDA_ARCHITECTURES ON)
+  # Torch/PyTorch propagates C-only warning flags (e.g. -Wno-duplicate-decl-specifier)
+  # via imported target interface options and HIP variables.
+  # See: https://github.com/pytorch/pytorch/pull/164552
+  # Strip them so they don't get passed to the C++ compiler.
+  foreach(_target torch torch_cpu torch_cuda torch_hip c10 c10_cuda c10_hip)
+    if(TARGET ${_target})
+      get_target_property(_opts ${_target} INTERFACE_COMPILE_OPTIONS)
+      if(_opts)
+        string(REGEX REPLACE "-Wno-duplicate-decl-specifier" "" _opts "${_opts}")
+        set_target_properties(${_target} PROPERTIES INTERFACE_COMPILE_OPTIONS "${_opts}")
+      endif()
+    endif()
+  endforeach()
+  # Also strip from HIP-specific CMake variables that torch may have polluted
+  foreach(_var CMAKE_CXX_FLAGS CMAKE_HIP_FLAGS HIP_CXX_FLAGS HIP_HIPCC_FLAGS)
+    if(DEFINED ${_var})
+      string(REPLACE "-Wno-duplicate-decl-specifier" "" ${_var} "${${_var}}")
+    endif()
+  endforeach()
+  list(APPEND AMS_APP_DEFINES "__AMS_ENABLE_TORCH__")
+else()
+  message(STATUS "PyTorch support disabled (ENABLE_TORCH=OFF). ML inference will not be available.")
+endif()
+
+# ------------------------------------------------------------------------------
+if (WITH_RZ)
+  find_package(MPI REQUIRED)
+  add_subdirectory(rz)
+  list(APPEND AMS_APP_INCLUDES "${RZ_AMS_INCLUDES}" "${MPI_INCLUDE_PATH}")
+  list(APPEND AMS_APP_LIB_DIRS "${RZ_AMS_LIBDIRS}")
+  list(APPEND AMS_APP_LIBRARIES "${RZ_AMS_LIBRARIES}" "${MPI_C_LIBRARIES}")
+  list(APPEND AMS_APP_DEFINES "${RZ_AMS_DEFINES}")
+endif()
 
 if (ENABLE_PERFFLOWASPECT)
   find_package(perfflowaspect CONFIG REQUIRED)
@@ -406,9 +453,12 @@ endif()
 if (NOT BUILD_SHARED_LIBS)
   # Pin the location hints (see ams_append_pinned_dependency above)
   ams_append_pinned_dependency("find_dependency(fmt CONFIG REQUIRED)" fmt_DIR)
-  ams_append_pinned_dependency("find_dependency(Torch REQUIRED)" Torch_DIR)
   ams_append_pinned_dependency("find_dependency(nlohmann_json REQUIRED)" nlohmann_json_DIR)
   ams_append_pinned_dependency("find_dependency(tl-expected REQUIRED)" tl-expected_DIR)
+
+  if (ENABLE_TORCH)
+    ams_append_pinned_dependency("find_dependency(Torch CONFIG REQUIRED)" Torch_DIR)
+  endif()
   if (ENABLE_MPI)
     ams_append_package_dependency("find_dependency(MPI REQUIRED COMPONENTS C CXX)")
   endif()
@@ -419,7 +469,7 @@ if (NOT BUILD_SHARED_LIBS)
     ams_append_pinned_dependency("find_dependency(HIP REQUIRED)" hip_DIR)
   endif()
   if (ENABLE_CALIPER)
-    ams_append_pinned_dependency("find_dependency(caliper REQUIRED)" caliper_DIR)
+    ams_append_pinned_dependency("find_dependency(caliper CONFIG REQUIRED)" caliper_DIR)
   endif()
   if (AMS_HDF5_MODE STREQUAL "HDF5_STATIC_TARGET")
     ams_append_pinned_dependency("find_dependency(HDF5 CONFIG REQUIRED COMPONENTS C static)" HDF5_DIR)

INSTALL.md

@@ -198,6 +198,7 @@ source scripts/gitlab/setup-env.sh
 
 cmake -S . -B build \
   -DCMAKE_BUILD_TYPE=Release \
+<<<<<<< HEAD
   -DBUILD_SHARED_LIBS=On \
   -DCMAKE_INSTALL_RPATH_USE_LINK_PATH=On \
   -DENABLE_MPI=On \
@@ -213,6 +214,20 @@ cmake -S . -B build \
   -DAMS_FMT_DIR="$AMS_FMT_DIR" \
   -Dnlohmann_json_DIR="$AMS_NLOHMANN_JSON_DIR" \
   -Dtl-expected_DIR="$AMS_TL_EXPECTED_DIR"
+=======
+  -DWITH_CUDA=On \
+  -DUMPIRE_DIR=$AMS_UMPIRE_PATH \
+  -DMFEM_DIR=$AMS_MFEM_PATH \
+  -DWITH_FAISS=On \
+  -DWITH_MPI=On \
+  -DENABLE_TORCH=On \
+  -DWITH_TESTS=Off \
+  -DTorch_DIR=$AMS_TORCH_PATH \
+  -DFAISS_DIR=$AMS_FAISS_PATH \
+  -DAMS_CUDA_ARCH=${AMS_CUDA_ARCH} \
+  -DWITH_AMS_DEBUG=On \
+  ../
+>>>>>>> 4598085 (WIP. Removing Torch from the main path (all tests green with WITH_TORCH=On except integration tests))
 
 cmake --build build -j 6
 cmake --install build

cmake/AMSConfig.cmake.in

@@ -7,6 +7,7 @@ find_dependency(Threads REQUIRED)
 set(AMS_ENABLE_MPI @ENABLE_MPI@)
 set(AMS_ENABLE_CUDA @ENABLE_CUDA@)
 set(AMS_ENABLE_HIP @ENABLE_HIP@)
+set(AMS_ENABLE_TORCH @ENABLE_TORCH@)
 set(AMS_ENABLE_CALIPER @ENABLE_CALIPER@)
 set(AMS_ENABLE_WORKFLOW @ENABLE_WORKFLOW@)
 set(AMS_ENABLE_RMQ @ENABLE_RMQ@)