Merge branch 'elixir-nx:main' into docs_refact

Author: TomasPegado

exla/Makefile

@@ -22,10 +22,8 @@ XLA_EXTENSION_LIB_LINK_PATH = ../$(CWD_RELATIVE_TO_PRIV_PATH)/$(XLA_EXTENSION_LI
 EXLA_CACHE_SO_LINK_PATH = $(CWD_RELATIVE_TO_PRIV_PATH)/$(EXLA_CACHE_SO)
 
 # Build flags
-# c++17 is needed, otherwise xla headers
-# break on some conflicting llvm/std definitions
-# Note: this is on :xla 0.5.0 -- things can change with later versions
-CFLAGS = -fPIC -I$(ERTS_INCLUDE_DIR) -I$(XLA_INCLUDE_PATH) -Wall -Wno-sign-compare \
+# Note that XLA requires c++17, Fine as well
+CFLAGS = -fPIC -I$(ERTS_INCLUDE_DIR) -I$(FINE_INCLUDE_DIR) -I$(XLA_INCLUDE_PATH) -Wall -Wno-sign-compare \
 	 -Wno-unused-parameter -Wno-missing-field-initializers -Wno-comment \
 	 -std=c++17 -w -DLLVM_VERSION_STRING=
 
@@ -38,27 +36,29 @@ else
 	CFLAGS += -O3
 endif
 
-LDFLAGS = -L$(XLA_EXTENSION_LIB) -lxla_extension -shared
+NVCC := $(CXX)
+NVCCFLAGS = $(CFLAGS)
+LDFLAGS = -L$(XLA_EXTENSION_LIB) -lxla_extension -shared -fvisibility=hidden
 
 ifeq ($(CROSSCOMPILE),)
 	# Interrogate the system for local compilation
 	UNAME_S = $(shell uname -s)
 
+ifdef ($(EXLA_CPU_ONLY),)
+$(info EXLA_CPU_ONLY is not set, checking for nvcc availability)
 	NVCC_RESULT := $(shell which nvcc 2> /dev/null)
 	NVCC_TEST := $(notdir $(NVCC_RESULT))
 
-ifeq ($(NVCC_TEST),nvcc)
-	NVCC := nvcc
-	NVCCFLAGS += -DCUDA_ENABLED
+	ifeq ($(NVCC_TEST),nvcc)
+		NVCC := nvcc
+		NVCCFLAGS += -DCUDA_ENABLED
+	endif
 else
-	NVCC := $(CXX)
-	NVCCFLAGS = $(CFLAGS)
+$(info EXLA_CPU_ONLY is set, skipping nvcc step)
 endif
 else
 	# Determine settings for cross-compiled builds like for Nerves
 	UNAME_S = Linux
-	NVCC := $(CXX)
-	NVCCFLAGS = $(CFLAGS)
 endif
 
 ifeq ($(UNAME_S), Darwin)
@@ -82,7 +82,7 @@ $(EXLA_SO): $(EXLA_CACHE_SO)
 		ln -sf $(EXLA_CACHE_SO_LINK_PATH) $(EXLA_SO) ; \
 	fi
 
-SOURCES = $(EXLA_DIR)/exla.cc $(EXLA_DIR)/exla_client.cc $(EXLA_DIR)/exla_mlir.cc $(EXLA_DIR)/custom_calls.cc $(EXLA_DIR)/exla_nif_util.cc $(EXLA_DIR)/ipc.cc
+SOURCES = $(EXLA_DIR)/exla.cc $(EXLA_DIR)/exla_client.cc $(EXLA_DIR)/exla_mlir.cc $(EXLA_DIR)/custom_calls.cc $(EXLA_DIR)/ipc.cc
 SOURCES += $(wildcard $(EXLA_DIR)/custom_calls/*.cc)
 HEADERS = $(EXLA_DIR)/exla_mlir.h $(EXLA_DIR)/custom_calls/qr.h $(EXLA_DIR)/custom_calls/eigh.h $(EXLA_DIR)/exla_client.h $(EXLA_DIR)/exla_nif_util.h $(EXLA_DIR)/exla_log_sink.h $(EXLA_DIR)/ipc.h
 OBJECTS = $(patsubst $(EXLA_DIR)/%.cc,$(EXLA_CACHE_OBJ_DIR)/%.o,$(SOURCES)) $(EXLA_CACHE_OBJ_DIR)/exla_cuda.o

exla/README.md

@@ -43,6 +43,8 @@ EXLA relies on the [XLA](https://github.com/elixir-nx/xla) package to provide th
   * Incompatible protocol buffer versions
     * Error message: "this file was generated by an older version of protoc which is incompatible with your Protocol Buffer headers".
     * If you have `protoc` installed on your machine, it may conflict with the `protoc` precompiled inside XLA. Uninstall, unlink, or remove `protoc` from your path to continue.
+  * Missing CUDA symbols
+    * In some cases, you might be compiling a CPU-only version of `:xla` in an environment that has CUDA available. For these cases, you can set the `EXLA_CPU_ONLY` environment variable to any value to disable custom CUDA functionality in EXLA.
 
 ### Usage with Nerves
 

exla/c_src/exla/custom_calls/eigh.h

@@ -2,11 +2,18 @@
 
 #include "Eigen/Eigenvalues"
 
+#include <algorithm>
 #include <iostream>
+#include <numeric>
+#include <vector>
 
 template <typename DataType>
-void single_matrix_eigh_cpu_custom_call(DataType *eigenvalues_out, DataType *eigenvectors_out, DataType *in, uint64_t m, uint64_t n) {
-  typedef Eigen::Matrix<DataType, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> RowMajorMatrix;
+void single_matrix_eigh_cpu_custom_call(DataType *eigenvalues_out,
+                                        DataType *eigenvectors_out,
+                                        DataType *in, uint64_t m, uint64_t n) {
+  typedef Eigen::Matrix<DataType, Eigen::Dynamic, Eigen::Dynamic,
+                        Eigen::RowMajor>
+      RowMajorMatrix;
 
   // Map the input matrix
   Eigen::Map<RowMajorMatrix> input(in, m, n);
@@ -20,14 +27,32 @@ void single_matrix_eigh_cpu_custom_call(DataType *eigenvalues_out, DataType *eig
   }
 
   // Get the eigenvalues and eigenvectors
-  Eigen::Matrix<DataType, Eigen::Dynamic, 1> eigenvalues = eigensolver.eigenvalues();
+  Eigen::Matrix<DataType, Eigen::Dynamic, 1> eigenvalues =
+      eigensolver.eigenvalues();
   RowMajorMatrix eigenvectors = eigensolver.eigenvectors();
 
-  // Copy the eigenvalues to the output
-  std::memcpy(eigenvalues_out, eigenvalues.data(), m * sizeof(DataType));
+  // Create a vector of indices and sort it based on eigenvalues in decreasing
+  // order
+  std::vector<int> indices(m);
+  std::iota(indices.begin(), indices.end(), 0);
+  std::sort(indices.begin(), indices.end(), [&eigenvalues](int i, int j) {
+    return std::abs(eigenvalues(i)) > std::abs(eigenvalues(j));
+  });
+
+  // Sort eigenvalues and rearrange eigenvectors
+  Eigen::Matrix<DataType, Eigen::Dynamic, 1> sorted_eigenvalues(m);
+  RowMajorMatrix sorted_eigenvectors(m, n);
+  for (int i = 0; i < m; ++i) {
+    sorted_eigenvalues(i) = eigenvalues(indices[i]);
+    sorted_eigenvectors.col(i) = eigenvectors.col(indices[i]);
+  }
+
+  // Copy the sorted eigenvalues to the output
+  std::memcpy(eigenvalues_out, sorted_eigenvalues.data(), m * sizeof(DataType));
 
-  // Copy the eigenvectors to the output
-  std::memcpy(eigenvectors_out, eigenvectors.data(), m * n * sizeof(DataType));
+  // Copy the sorted eigenvectors to the output
+  std::memcpy(eigenvectors_out, sorted_eigenvectors.data(),
+              m * n * sizeof(DataType));
 }
 
 template <typename DataType>
@@ -40,18 +65,22 @@ void eigh_cpu_custom_call(void *out[], const void *in[]) {
   uint64_t num_eigenvectors_dims = dim_sizes[2];
 
   uint64_t *operand_dims_ptr = (uint64_t *)in[2];
-  std::vector<uint64_t> operand_dims(operand_dims_ptr, operand_dims_ptr + num_operand_dims);
+  std::vector<uint64_t> operand_dims(operand_dims_ptr,
+                                     operand_dims_ptr + num_operand_dims);
 
   uint64_t *eigenvalues_dims_ptr = (uint64_t *)in[3];
-  std::vector<uint64_t> eigenvalues_dims(eigenvalues_dims_ptr, eigenvalues_dims_ptr + num_eigenvalues_dims);
+  std::vector<uint64_t> eigenvalues_dims(
+      eigenvalues_dims_ptr, eigenvalues_dims_ptr + num_eigenvalues_dims);
 
   uint64_t *eigenvectors_dims_ptr = (uint64_t *)in[4];
-  std::vector<uint64_t> eigenvectors_dims(eigenvectors_dims_ptr, eigenvectors_dims_ptr + num_eigenvectors_dims);
+  std::vector<uint64_t> eigenvectors_dims(
+      eigenvectors_dims_ptr, eigenvectors_dims_ptr + num_eigenvectors_dims);
 
   uint64_t m = eigenvectors_dims[eigenvectors_dims.size() - 2];
   uint64_t n = eigenvectors_dims[eigenvectors_dims.size() - 1];
 
-  auto leading_dimensions = std::vector<uint64_t>(operand_dims.begin(), operand_dims.end() - 2);
+  auto leading_dimensions =
+      std::vector<uint64_t>(operand_dims.begin(), operand_dims.end() - 2);
 
   uint64_t batch_items = 1;
   for (uint64_t i = 0; i < leading_dimensions.size(); i++) {
@@ -61,15 +90,16 @@ void eigh_cpu_custom_call(void *out[], const void *in[]) {
   DataType *eigenvalues = (DataType *)out[0];
   DataType *eigenvectors = (DataType *)out[1];
 
-  uint64_t eigenvalues_stride = eigenvalues_dims[eigenvalues_dims.size() - 1] * sizeof(DataType);
-  uint64_t eigenvectors_stride = eigenvectors_dims[eigenvectors_dims.size() - 1] * eigenvectors_dims[eigenvectors_dims.size() - 2] * sizeof(DataType);
-  uint64_t inner_stride = m * n * sizeof(DataType);
+  uint64_t eigenvalues_stride = eigenvalues_dims[eigenvalues_dims.size() - 1];
+  uint64_t eigenvectors_stride =
+      eigenvectors_dims[eigenvectors_dims.size() - 1] *
+      eigenvectors_dims[eigenvectors_dims.size() - 2];
+  uint64_t inner_stride = m * n;
 
   for (uint64_t i = 0; i < batch_items; i++) {
     single_matrix_eigh_cpu_custom_call<DataType>(
         eigenvalues + i * eigenvalues_stride,
-        eigenvectors + i * eigenvectors_stride,
-        operand + i * inner_stride / sizeof(DataType),
-        m, n);
+        eigenvectors + i * eigenvectors_stride, operand + i * inner_stride, m,
+        n);
   }
 }