using :cmake compiler in mix.exs

torchx/CMakeLists.txt

@@ -24,12 +24,26 @@ if(NOT IS_DIRECTORY "${PRIV_DIR}")
   file(MAKE_DIRECTORY "${PRIV_DIR}")
 endif()
 
+if(DEFINED ERTS_INCLUDE_DIR AND NOT "${ERTS_INCLUDE_DIR}" STREQUAL "")
+    set(ERTS_INCLUDE_DIR "${ERTS_INCLUDE_DIR}")
+else()
+    set(ERTS_INCLUDE_DIR_ONE_LINER "erl -noshell -eval \"io:format('~ts/erts-~ts/include/', [code:root_dir(), erlang:system_info(version)]), halt().\"")
+    if(WIN32)
+        execute_process(COMMAND powershell -command "${ERTS_INCLUDE_DIR_ONE_LINER}" OUTPUT_VARIABLE ERTS_INCLUDE_DIR)
+    else()
+        execute_process(COMMAND bash -c "${ERTS_INCLUDE_DIR_ONE_LINER}" OUTPUT_VARIABLE ERTS_INCLUDE_DIR)
+    endif()
+    set(ERTS_INCLUDE_DIR "${ERTS_INCLUDE_DIR}")
+endif()
+include_directories(${ERTS_INCLUDE_DIR})
+
 set(C_SRC "${CMAKE_CURRENT_SOURCE_DIR}/c_src")
 set(Torch_DIR "$ENV{LIBTORCH_DIR}/share/cmake/Torch")
 
 message(STATUS "C_SRC: ${C_SRC}")
 message(STATUS "MIX_APP_PATH: $ENV{MIX_APP_PATH}")
 message(STATUS "PRIV_DIR: ${PRIV_DIR}")
+message(STATUS "ERTS_INCLUDE_DIR: ${ERTS_INCLUDE_DIR}")
 message(STATUS "LIBTORCH_DIR: $ENV{LIBTORCH_DIR}")
 message(STATUS "LIBTORCH_BASE: $ENV{LIBTORCH_BASE}")
 message(STATUS "MIX_BUILD_EMBEDDED $ENV{MIX_BUILD_EMBEDDED}")
@@ -110,7 +124,7 @@ if(MSVC)
                     ${TORCH_DLLS}
                     "${PRIV_DIR}/$ENV{LIBTORCH_BASE}")
 else()
-  if($ENV{MIX_BUILD_EMBEDDED} STREQUAL "true")
+  if(ENV{MIX_BUILD_EMBEDDED} STREQUAL "true")
     set(EMBEDDED_TYPE copy_directory)
   else()
     set(EMBEDDED_TYPE create_symlink)

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,53 @@ 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()
+    priv_path = Path.join(Mix.Project.app_path(), "priv")
+    libtorch_link_path =
+      libtorch_config.env_dir || relative_to(libtorch_config.dir, priv_path)
+
+    System.put_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.Project.app_path()
+    })
+
+    cmd!(cmake, ["-S", ".", "-B", cmake_build_dir])
+    cmd!(cmake, ["--build", cmake_build_dir, "--config", cmake_build_type])
+    cmd!(cmake, ["--install", cmake_build_dir, "--config", cmake_build_type])
+
+    {:ok, []}
+  end
+
+  defp cmd!(exec, args) do
+    opts = [
+      into: IO.stream(:stdio, :line),
+      stderr_to_stdout: true,
+    ]
+
+    case System.cmd(exec, args, opts) do
+      {_, 0} -> :ok
+      {_, status} -> Mix.raise("cmake failed with status #{status}")
+    end
+  end
 end