Adding support for ADD_TORCH_LIBRARY macro

3 files changed, 37 insertions, 21 deletions

diff --git a/cmake/TorchPackage.cmake b/cmake/TorchPackage.cmake
index 7fcbdff..f966dac 100644
--- a/cmake/TorchPackage.cmake
+++ b/cmake/TorchPackage.cmake
@@ -1,5 +1,21 @@
 # -*- cmake -*-
 
+MACRO(ADD_TORCH_LIBRARY package type src)
+  IF ("${type}" STREQUAL "STATIC")
+    if ("${src}" MATCHES "cu$" OR "${src}" MATCHES "cu;")
+      CUDA_ADD_LIBRARY(${package} STATIC ${src})
+    else()
+      ADD_LIBRARY(${package} STATIC ${src})
+    endif()
+  ELSE()
+    if ("${src}" MATCHES "cu$" OR "${src}" MATCHES "cu;")
+      CUDA_ADD_LIBRARY(${package} ${type} ${src})
+    else()
+      ADD_LIBRARY(${package} ${type} ${src})
+    endif()
+  ENDIF()
+ENDMACRO()
+
 MACRO(ADD_TORCH_PACKAGE package src luasrc)
   INCLUDE_DIRECTORIES(${CMAKE_CURRENT_SOURCE_DIR})
   INCLUDE_DIRECTORIES(${Torch_LUA_INCLUDE_DIR})
@@ -8,17 +24,7 @@ MACRO(ADD_TORCH_PACKAGE package src luasrc)
  # As per CMake doc, macro arguments are not variables, so simple test syntax not working
   IF(NOT "${src}" STREQUAL "")
 
-    if ("${src}" MATCHES "cu$" OR "${src}" MATCHES "cu;")
-      CUDA_ADD_LIBRARY(${package} MODULE ${src})
-      if(BUILD_STATIC)
-        CUDA_ADD_LIBRARY(${package}_static STATIC ${src})
-      endif()
-    else()
-      ADD_LIBRARY(${package} MODULE ${src})
-      if(BUILD_STATIC)
-        ADD_LIBRARY(${package}_static STATIC ${src})
-      endif()
-    endif()
+    ADD_TORCH_LIBRARY(${package} MODULE "${src}")
 
     ### Torch packages supposes libraries prefix is "lib"
     SET_TARGET_PROPERTIES(${package} PROPERTIES
@@ -31,12 +37,13 @@ MACRO(ADD_TORCH_PACKAGE package src luasrc)
         LINK_FLAGS "-undefined dynamic_lookup")
     ENDIF()
 
-    if(BUILD_STATIC)
+    IF (BUILD_STATIC OR "$ENV{STATIC_TH}" STREQUAL "YES")
+      ADD_TORCH_LIBRARY(${package}_static STATIC "${src}")
       SET_TARGET_PROPERTIES(${package}_static PROPERTIES
         COMPILE_FLAGS "-fPIC")
       SET_TARGET_PROPERTIES(${package}_static PROPERTIES
         PREFIX "lib" IMPORT_PREFIX "lib" OUTPUT_NAME "${package}")
-    endif()
+    ENDIF()
 
     INSTALL(TARGETS ${package}
       RUNTIME DESTINATION ${Torch_INSTALL_LUA_CPATH_SUBDIR}
diff --git a/lib/TH/CMakeLists.txt b/lib/TH/CMakeLists.txt
index c4e6694..b481a8a 100644
--- a/lib/TH/CMakeLists.txt
+++ b/lib/TH/CMakeLists.txt
@@ -145,7 +145,6 @@ IF(C_AVX2_FOUND)
   SET(CMAKE_C_FLAGS "-DUSE_AVX2 ${CMAKE_C_FLAGS}")
 ENDIF(C_AVX2_FOUND)
 
-
 CHECK_C_SOURCE_RUNS("
 #include <stdatomic.h>
 int main()
@@ -243,10 +242,15 @@ SET(src ${src} ${hdr} ${simd})
 ##### build section
 ######################################################################
 
-ADD_LIBRARY(TH SHARED ${src})
-if(BUILD_STATIC)
-  ADD_LIBRARY(TH_static STATIC ${src})
-endif()
+ADD_TORCH_LIBRARY(TH SHARED "${src}")
+
+IF (BUILD_STATIC OR "$ENV{STATIC_TH}" STREQUAL "YES")
+  ADD_TORCH_LIBRARY(TH_static STATIC "${src}")
+  SET_TARGET_PROPERTIES(TH_static PROPERTIES
+    COMPILE_FLAGS "-fPIC")
+  SET_TARGET_PROPERTIES(TH_static PROPERTIES
+    PREFIX "lib" IMPORT_PREFIX "lib" OUTPUT_NAME "TH")
+ENDIF()
 
 IF(NOT TH_SO_VERSION)
   SET(TH_SO_VERSION 0)
diff --git a/lib/luaT/CMakeLists.txt b/lib/luaT/CMakeLists.txt
index f33768c..072991c 100644
--- a/lib/luaT/CMakeLists.txt
+++ b/lib/luaT/CMakeLists.txt
@@ -9,9 +9,14 @@ IF(LUALIB)
 ENDIF()
 
 ADD_LIBRARY(luaT SHARED luaT.h luaT.c)
-if(BUILD_STATIC)
+
+IF (BUILD_STATIC OR "$ENV{STATIC_TH}" STREQUAL "YES")
   ADD_LIBRARY(luaT_static STATIC luaT.h luaT.c)
-endif()
+  SET_TARGET_PROPERTIES(luaT_static PROPERTIES
+    COMPILE_FLAGS "-fPIC")
+  SET_TARGET_PROPERTIES(luaT_static PROPERTIES
+    PREFIX "lib" IMPORT_PREFIX "lib" OUTPUT_NAME "luaT")
+ENDIF()
 
 SET_TARGET_PROPERTIES(luaT PROPERTIES
   VERSION   0
@@ -41,5 +46,5 @@ GET_FILENAME_COMPONENT(LUAT_OUTPUT_NAME ${LUAT_OUTPUT_NAME} NAME)
 SET(LUAT_LIBRARIES "${Torch_INSTALL_LIB}/${LUAT_OUTPUT_NAME}")
 SET(LUAT_INCLUDE_DIR "${Torch_INSTALL_INCLUDE}")
 CONFIGURE_FILE(luaTConfig.cmake.in "${CMAKE_CURRENT_BINARY_DIR}/cmake-exports/luaTConfig.cmake")
-INSTALL(FILES "${CMAKE_CURRENT_BINARY_DIR}/cmake-exports/luaTConfig.cmake" 
+INSTALL(FILES "${CMAKE_CURRENT_BINARY_DIR}/cmake-exports/luaTConfig.cmake"
   DESTINATION "${Torch_INSTALL_CMAKE_SUBDIR}")