Add linker args for more torch libs

.

Author: Iainmon

CMakeLists.txt

@@ -48,20 +48,44 @@ target_include_directories(
 set(BRIDGE_OBJECT_FILES $<TARGET_OBJECTS:bridge>)
 
 
-file(GLOB PYTORCH_LIBS "${LIBTORCH_DIR}/lib/*.a" "${LIBTORCH_DIR}/lib/*.dylib" "${LIBTORCH_DIR}/lib/*.so")
-
-set(PYTORCH_LIBS_LINKER_ARGS "-ltorch")  # Will hold the list of "-l..." flags.
-foreach(lib_path IN LISTS PYTORCH_LIBS)
-  # Get just the filename without the directory or extension
-  get_filename_component(lib_name "${lib_path}" NAME_WE)
-  # If it starts with "lib", strip that off
-  string(REGEX REPLACE "^lib" "" lib_name "${lib_name}")
-  # Now prepend "-l" to the actual library name
-  list(APPEND PYTORCH_LIBS_LINKER_ARGS "-l${lib_name}")
+
+
+file(GLOB LIBTORCH_ALL_LIB_FILES "${LIBTORCH_DIR}/lib/*.a" "${LIBTORCH_DIR}/lib/*.dylib" "${LIBTORCH_DIR}/lib/*.so")
+
+set(LIBTORCH_ALL_LIBS "")
+foreach(lib_path IN LISTS LIBTORCH_ALL_LIB_FILES)
+    get_filename_component(lib_name "${lib_path}" NAME_WE)
+    list(APPEND LIBTORCH_ALL_LIBS "${lib_name}")
+endforeach()
+
+
+set(REQUIRED_LIBS
+    "libtorch"
+    "libtorch_cpu"
+    "libc10"
+    "libtorch_global_deps"
+)
+
+set(DISALLOWED_LIBS
+    "libtorch_python"
+)
+
+
+set(LIBTORCH_LIBS_LINKER_ARGS "")  # Will hold the list of "-l..." flags.
+foreach(lib_name IN LISTS LIBTORCH_ALL_LIBS)
+    if(lib_name IN_LIST DISALLOWED_LIBS)
+        if(lib_name IN_LIST REQUIRED_LIBS)
+            message(FATAL_ERROR "Required lib ${lib_name} is disallowed.")
+        else()
+            message(STATUS "Skipping disallowed lib: ${lib_name}")
+            continue()
+        endif()
+    endif()
+    string(REGEX REPLACE "^lib" "" lib_name_short "${lib_name}")
+    list(APPEND LIBTORCH_LIBS_LINKER_ARGS "-l${lib_name_short}")
 endforeach()
 
-# cmake_print_variables(PYTORCH_LIBS)
-# cmake_print_variables(PYTORCH_LIBS_LINKER_ARGS)
+cmake_print_variables(LIBTORCH_LIBS_LINKER_ARGS)
 
 add_executable(TorchBridge ${BRIDGE_DIR}/lib/Bridge.chpl)
 add_dependencies(TorchBridge bridge)
@@ -70,10 +94,11 @@ target_link_options(TorchBridge
         ${BRIDGE_DIR}/include/bridge.h
         ${BRIDGE_OBJECT_FILES}
         -L ${LIBTORCH_DIR}/lib
-        "-ltorch"
-        "-ltorch_cpu"
-        "-lc10"
-        # ${PYTORCH_LIBS_LINKER_ARGS}
+        # "-ltorch"
+        # "-ltorch_cpu"
+        # "-lc10"
+        # "-ltorch_global_deps"
+        ${LIBTORCH_LIBS_LINKER_ARGS}
         --ldflags "-Wl,-rpath,${LIBTORCH_DIR}/lib"
 )
 # install(TARGETS TorchBridge DESTINATION ".")

bridge/lib/Bridge.chpl

@@ -15,6 +15,7 @@ extern record bridge_tensor_t {
 }
 
 
+
 extern proc increment2(arr: [] real(32), sizes: [] int(32), dim: int(32)): bridge_tensor_t;
 extern proc increment3(in arr: bridge_tensor_t): bridge_tensor_t;
 

cmake/LibTorchDL.cmake

@@ -6,6 +6,9 @@ include(dlcache)
 
 include(CMakePrintHelpers)
 
+
+# set(TORCH_COMPUTE_PLATFORM "cpu") # cu118, cu124, cu126, cpu
+# set(TORCH_VERSION "2.1.0")
 set(TORCH_URL_PREFIX "https://download.pytorch.org/libtorch/nightly/cpu")
 
 
@@ -20,7 +23,7 @@ endif()
 if(APPLE)
     set(TORCH_DISTRIBUTION "${TORCH_URL_PREFIX}/libtorch-macos-arm64-latest.zip")
 elseif(LINUX)
-    set(TORCH_DISTRIBUTION "$${TORCH_URL_PREFIX}/libtorch-cxx11-abi-static-with-deps-latest.zip")
+    set(TORCH_DISTRIBUTION "$${TORCH_URL_PREFIX}/libtorch-shared-with-deps-latest.zip")
 endif()
 
 function(download_libtorch)