Fix exla and libtorch for macos (#463)

elixir-nx · Sep 19, 2021 · 432bc74 · dognotdog · Sep 20, 2021 · 432bc74
1 parent 6e2e19a
commit 432bc74
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 16 deletions.
diff --git a/exla/Makefile b/exla/Makefile
@@ -1,15 +1,15 @@
 # Environment variables passed via elixir_make
 # ERTS_INCLUDE_DIR
+# MIX_APP_PATH
 
 # XLA Extension Installation Location
 XLA_EXTENSION_DIR ?= cache/xla_extension
-
 XLA_EXTENSION_LIB = $(XLA_EXTENSION_DIR)/lib
 XLA_INCLUDE_PATH = $(XLA_EXTENSION_DIR)/include
 
 # Private configuration
 EXLA_DIR = c_src/exla
-PRIV_DIR = priv
+PRIV_DIR = $(MIX_APP_PATH)/priv
 EXLA_SO = $(PRIV_DIR)/libexla.so
 EXLA_LIB_DIR = $(PRIV_DIR)/lib
 
@@ -18,19 +18,25 @@ CFLAGS = -fPIC -I$(ERTS_INCLUDE_DIR) -isystem $(XLA_INCLUDE_PATH) -O3 -Wall -Wex
 				 -Wno-unused-parameter -Wno-missing-field-initializers -Wno-comment \
 				 -shared -std=c++14
 
-# Use a relative RPATH, so at runtime libexla.so looks for libxla_extension.so
-# in ./lib regardless of the absolute location. This way priv can be safely
-# packed into an Elixir release. Also, we use $$ to escape Makefile variable
-# and single quotes to escape shell variable
-LDFLAGS = -L$(EXLA_LIB_DIR) -Wl,-rpath,'$$ORIGIN/lib' -lxla_extension
+LDFLAGS = -L$(EXLA_LIB_DIR) -lxla_extension
+
+ifeq ($(shell uname -s), Darwin)
+	LDFLAGS += -flat_namespace -undefined suppress
+	POST_INSTALL = install_name_tool -change bazel-out/darwin-opt/bin/tensorflow/compiler/xla/extension/libxla_extension.so @loader_path/lib/libxla_extension.so $(EXLA_SO)
+else
+	# Use a relative RPATH, so at runtime libexla.so looks for libxla_extension.so
+	# in ./lib regardless of the absolute location. This way priv can be safely
+	# packed into an Elixir release. Also, we use $$ to escape Makefile variable
+	# and single quotes to escape shell variable
+	LDFLAGS += -Wl,-rpath,'$$ORIGIN/lib'
+	POST_INSTALL = $(NOOP)
+endif
 
 $(EXLA_SO): $(XLA_EXTENSION_DIR) $(EXLA_DIR)/exla.cc $(EXLA_DIR)/exla_client.cc $(EXLA_DIR)/exla_client.h $(EXLA_DIR)/exla_nif_util.cc $(EXLA_DIR)/exla_nif_util.h $(EXLA_DIR)/exla_log_sink.h
-	mkdir -p $(PRIV_DIR) && \
-		ln -sf $(abspath $(XLA_EXTENSION_LIB)) $(EXLA_LIB_DIR) && \
-		$(CXX) $(CFLAGS) $(EXLA_DIR)/exla.cc $(EXLA_DIR)/exla_nif_util.h \
-											$(EXLA_DIR)/exla_nif_util.cc $(EXLA_DIR)/exla_client.h \
-											$(EXLA_DIR)/exla_client.cc $(EXLA_DIR)/exla_log_sink.h \
-											-o $(EXLA_SO) $(LDFLAGS)
+	mkdir -p $(PRIV_DIR)
+	ln -sf $(abspath $(XLA_EXTENSION_LIB)) $(EXLA_LIB_DIR)
+	$(CXX) $(CFLAGS) $(EXLA_DIR)/exla.cc $(EXLA_DIR)/exla_nif_util.cc $(EXLA_DIR)/exla_client.cc -o $(EXLA_SO) $(LDFLAGS)
+	$(POST_INSTALL)
 
 clean:
 	rm -rf $(EXLA_SO)
diff --git a/torchx/Makefile b/torchx/Makefile
@@ -1,5 +1,5 @@
 LIBTORCH_VERSION := $(shell head -c 3 $(LIBTORCH_DIR)/build-version)
-MINIMUM_LIBTORCH_VERSION := 1.8
+EXPECTED_LIBTORCH_VERSION := 1.8
 
 LIBTORCH_INCLUDE_PATH = $(LIBTORCH_DIR)/include/
 LIBTORCH_API_INCLUDE_PATH = $(LIBTORCH_DIR)/include/torch/csrc/api/include
@@ -26,8 +26,8 @@ check_version:
 		exit 1; \
 	fi
 
-	@ if [ $(LIBTORCH_VERSION) >= $(MINIMUM_LIBTORCH_VERSION) ]; then \
-		echo "LibTorch version should be at least $(MINIMUM_LIBTORCH_VERSION). Got: $(LIBTORCH_VERSION)."; \
+	@ if [ $(LIBTORCH_VERSION) != $(EXPECTED_LIBTORCH_VERSION) ]; then \
+		echo "LibTorch version should be $(EXPECTED_LIBTORCH_VERSION). Got: $(LIBTORCH_VERSION)."; \
 		echo "You can download LibTorch here: https://pytorch.org/get-started/locally/"; \
 		exit 1; \
 	fi