Convert Makefile to CMakeLists.txt (#1389)

Co-authored-by: José Valim <[email protected]>
elixir-nx · Dec 5, 2023 · 55f745e · 55f745e
1 parent ad0cd2e
commit 55f745e
Show file tree

Hide file tree

Showing 4 changed files with 135 additions and 90 deletions.
diff --git a/torchx/CMakeLists.txt b/torchx/CMakeLists.txt
@@ -6,39 +6,68 @@ if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
   set(CMAKE_CUDA_ARCHITECTURES 75)
 endif()
 
+if(DEFINED ENV{CMAKE_TOOLCHAIN_FILE})
+  set(CMAKE_TOOLCHAIN_FILE $ENV{CMAKE_TOOLCHAIN_FILE})
+endif()
+
+if(NOT DEFINED CMAKE_BUILD_TYPE OR "${CMAKE_BUILD_TYPE}" STREQUAL "")
+    set(CMAKE_BUILD_TYPE "Release")
+endif()
+
+set(C_SRC "${CMAKE_CURRENT_SOURCE_DIR}/c_src")
+set(CMAKE_INSTALL_PREFIX "$ENV{PRIV_DIR}")
+if(NOT IS_DIRECTORY "$ENV{PRIV_DIR}")
+  file(MAKE_DIRECTORY "$ENV{PRIV_DIR}")
+endif()
+
+set(Torch_DIR "$ENV{LIBTORCH_DIR}/share/cmake/Torch")
+list(APPEND CMAKE_PREFIX_PATH $ENV{LIBTORCH_DIR})
+
+message(STATUS "CMAKE_TOOLCHAIN_FILE: $ENV{CMAKE_TOOLCHAIN_FILE}")
+message(STATUS "C_SRC: ${C_SRC}")
+message(STATUS "PRIV_DIR: $ENV{PRIV_DIR}")
+message(STATUS "LIBTORCH_DIR: $ENV{LIBTORCH_DIR}")
+message(STATUS "ERTS_INCLUDE_DIR: $ENV{ERTS_INCLUDE_DIR}")
+message(STATUS "LIBTORCH_BASE: $ENV{LIBTORCH_BASE}")
+message(STATUS "MIX_BUILD_EMBEDDED $ENV{MIX_BUILD_EMBEDDED}")
+message(STATUS "LIBTORCH_LINK $ENV{LIBTORCH_LINK}")
+
 find_package(Torch REQUIRED)
 
 STRING(REGEX REPLACE "\\\\" "/" C_SRC ${C_SRC})
 
 file(GLOB torchx_sources CONFIGURE_DEPENDS "${C_SRC}/*.cpp" "${C_SRC}/*.hpp")
 
-if (NOT APPLE)
 add_library(torchx SHARED ${torchx_sources})
-else()
-add_library(torchx MODULE ${torchx_sources})
-endif()
 
 target_link_libraries(torchx "${TORCH_LIBRARIES}")
-set_property(TARGET torchx PROPERTY CXX_STANDARD 14)
+set_property(TARGET torchx PROPERTY CXX_STANDARD 17)
 
-target_include_directories(torchx PUBLIC ${ERTS_INCLUDE_DIR})
+target_include_directories(torchx PUBLIC $ENV{ERTS_INCLUDE_DIR})
 
+install(
+    TARGETS torchx
+    DESTINATION "$ENV{PRIV_DIR}"
+)
 SET_TARGET_PROPERTIES(torchx PROPERTIES PREFIX "")
+if(NOT WIN32)
+    set_target_properties(torchx PROPERTIES SUFFIX ".so")
+endif()
 
 if (Torch_VERSION_MAJOR EQUAL 1)
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DUSING_TORCH_V1")
 endif()
 
 if(UNIX)
-set_target_properties(torchx PROPERTIES
-    INSTALL_RPATH_USE_LINK_PATH TRUE
-    BUILD_WITH_INSTALL_RPATH TRUE
-)
+  set_target_properties(torchx PROPERTIES
+      INSTALL_RPATH_USE_LINK_PATH TRUE
+      BUILD_WITH_INSTALL_RPATH TRUE
+  )
 
-if(NOT APPLE)
+  if(NOT APPLE)
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -shared")
-    set_target_properties(torchx PROPERTIES INSTALL_RPATH "\$ORIGIN/${LIBTORCH_BASE}")
-else()
+    set_target_properties(torchx PROPERTIES INSTALL_RPATH "\$ORIGIN/$ENV{LIBTORCH_BASE}")
+  else()
     # Although the compiler complains about not using these,
     # things only work with them set
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -undefined dynamic_lookup")
@@ -47,12 +76,39 @@ else()
         set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DMAC_ARM64")
     endif()
     # set(CMAKE_SHARED_LINKER_FLAGS "-bundle -flat_namespace -undefined suppress")
-    set_target_properties(torchx PROPERTIES INSTALL_RPATH "@loader_path/${LIBTORCH_BASE}")
-endif()
+    set_target_properties(torchx PROPERTIES INSTALL_RPATH "@loader_path/$ENV{LIBTORCH_BASE}")
+  endif()
 
-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fPIC -O3 -Wall -Wextra -Wno-unused-parameter -Wno-missing-field-initializers")
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fPIC -O3 -Wall -Wextra -Wno-unused-parameter -Wno-missing-field-initializers")
 else()
-# On Windows
-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /wd4624")
+  # On Windows
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /wd4624")
 endif()
 
+# https://pytorch.org/cppdocs/installing.html
+#
+# The following code block is suggested to be used on Windows.
+# According to https://github.com/pytorch/pytorch/issues/25457,
+# the DLLs need to be copied to avoid memory errors.
+if(MSVC)
+  file(GLOB TORCH_DLLS "$ENV{LIBTORCH_DIR}/lib/*.dll")
+  if(NOT IS_DIRECTORY "$ENV{PRIV_DIR}/$ENV{LIBTORCH_BASE}")
+   file(MAKE_DIRECTORY "$ENV{PRIV_DIR}/$ENV{LIBTORCH_BASE}")
+  endif()
+  add_custom_command(TARGET torchx
+                    POST_BUILD
+                    COMMAND ${CMAKE_COMMAND} -E copy_if_different
+                    ${TORCH_DLLS}
+                    "$ENV{PRIV_DIR}/$ENV{LIBTORCH_BASE}")
+else()
+  if(ENV{MIX_BUILD_EMBEDDED} STREQUAL "true")
+    set(EMBEDDED_TYPE copy_directory)
+  else()
+    set(EMBEDDED_TYPE create_symlink)
+  endif()
+  add_custom_command(TARGET torchx
+                      POST_BUILD
+                      COMMAND ${CMAKE_COMMAND} -E "${EMBEDDED_TYPE}"
+                      "$ENV{LIBTORCH_DIR}/lib"
+                      "$ENV{PRIV_DIR}/$ENV{LIBTORCH_BASE}")
+endif()
diff --git a/torchx/Makefile b/torchx/Makefile
diff --git a/torchx/Makefile.win b/torchx/Makefile.win
diff --git a/torchx/mix.exs b/torchx/mix.exs
@@ -4,7 +4,7 @@ defmodule Torchx.MixProject do
   @source_url "https://github.com/elixir-nx/nx"
   @version "0.7.0-dev"
 
-  @libtorch_compilers [:torchx, :elixir_make]
+  @libtorch_compilers [:torchx, :cmake]
 
   def project do
     [
@@ -26,22 +26,7 @@ defmodule Torchx.MixProject do
 
       # Compilers
       compilers: @libtorch_compilers ++ Mix.compilers(),
-      aliases: aliases(),
-      make_env: fn ->
-        libtorch_config = libtorch_config()
-
-        priv_path = Path.join(Mix.Project.app_path(), "priv")
-
-        libtorch_link_path =
-          libtorch_config.env_dir || relative_to(libtorch_config.dir, priv_path)
-
-        %{
-          "LIBTORCH_DIR" => libtorch_config.dir,
-          "LIBTORCH_BASE" => libtorch_config.base,
-          "MIX_BUILD_EMBEDDED" => "#{Mix.Project.config()[:build_embedded]}",
-          "LIBTORCH_LINK" => "#{libtorch_link_path}/lib"
-        }
-      end
+      aliases: aliases()
     ]
   end
 
@@ -95,7 +80,8 @@ defmodule Torchx.MixProject do
 
   defp aliases do
     [
-      "compile.torchx": &download_and_unzip/1
+      "compile.torchx": &download_and_unzip/1,
+      "compile.cmake": &cmake/1
     ]
   end
 
@@ -252,4 +238,61 @@ defmodule Torchx.MixProject do
 
   defp drop_common_prefix([h | left], [h | right]), do: drop_common_prefix(left, right)
   defp drop_common_prefix(left, right), do: {left, right}
+
+  def cmake(_args) do
+    priv? = File.dir?("priv")
+    Mix.Project.ensure_structure()
+
+    cmake = System.find_executable("cmake") || Mix.raise("cmake not found in the path")
+    cmake_build_type = System.get_env("CMAKE_BUILD_TYPE", "Release")
+    cmake_build_dir = Path.join(Mix.Project.app_path(), "cmake")
+    File.mkdir_p!(cmake_build_dir)
+
+    # IF there was no priv before and now there is one, we assume
+    # the user wants to copy it. If priv already existed and was
+    # written to it, then it won't be copied if build_embedded is
+    # set to true.
+    if not priv? and File.dir?("priv") do
+      Mix.Project.build_structure()
+    end
+
+    libtorch_config = libtorch_config()
+    mix_app_path = Mix.Project.app_path()
+    priv_path = Path.join(mix_app_path, "priv")
+
+    libtorch_link_path =
+      libtorch_config.env_dir || relative_to(libtorch_config.dir, priv_path)
+
+    erts_include_dir =
+      Path.join([:code.root_dir(), "erts-#{:erlang.system_info(:version)}", "include"])
+
+    env = %{
+      "LIBTORCH_DIR" => libtorch_config.dir,
+      "LIBTORCH_BASE" => libtorch_config.base,
+      "MIX_BUILD_EMBEDDED" => "#{Mix.Project.config()[:build_embedded]}",
+      "LIBTORCH_LINK" => "#{libtorch_link_path}/lib",
+      "MIX_APP_PATH" => mix_app_path,
+      "PRIV_DIR" => priv_path,
+      "ERTS_INCLUDE_DIR" => erts_include_dir
+    }
+
+    cmd!(cmake, ["-S", ".", "-B", cmake_build_dir], env)
+    cmd!(cmake, ["--build", cmake_build_dir, "--config", cmake_build_type], env)
+    cmd!(cmake, ["--install", cmake_build_dir, "--config", cmake_build_type], env)
+
+    {:ok, []}
+  end
+
+  defp cmd!(exec, args, env) do
+    opts = [
+      into: IO.stream(:stdio, :line),
+      stderr_to_stdout: true,
+      env: env
+    ]
+
+    case System.cmd(exec, args, opts) do
+      {_, 0} -> :ok
+      {_, status} -> Mix.raise("cmake failed with status #{status}")
+    end
+  end
 end