Update XLA

exla/Makefile

@@ -23,15 +23,15 @@ EXLA_CACHE_SO_LINK_PATH = $(CWD_RELATIVE_TO_PRIV_PATH)/$(EXLA_CACHE_SO)
 # Build flags
 CFLAGS = -fPIC -I$(ERTS_INCLUDE_DIR) -I$(XLA_INCLUDE_PATH) -O3 -Wall -Wno-sign-compare \
  -Wno-unused-parameter -Wno-missing-field-initializers -Wno-comment \
- -shared -std=c++17 -w -DLLVM_ON_UNIX=1
+ -shared -std=c++17 -w -DLLVM_VERSION_STRING=
 
 LDFLAGS = -L$(XLA_EXTENSION_LIB) -lxla_extension
 
 ifeq ($(shell uname -s), Darwin)
  LDFLAGS += -flat_namespace -undefined suppress
  POST_INSTALL = install_name_tool \
- -change bazel-out/darwin_arm64-opt/bin/tensorflow/compiler/xla/extension/libxla_extension.so @loader_path/xla_extension/lib/libxla_extension.so \
- -change bazel-out/darwin-opt/bin/tensorflow/compiler/xla/extension/libxla_extension.so @loader_path/xla_extension/lib/libxla_extension.so \
+ -change bazel-out/darwin_arm64-opt/bin/xla/extension/libxla_extension.so @loader_path/xla_extension/lib/libxla_extension.so \
+ -change bazel-out/darwin-opt/bin/xla/extension/libxla_extension.so @loader_path/xla_extension/lib/libxla_extension.so \
  $(EXLA_CACHE_SO)
 else
  # Use a relative RPATH, so at runtime libexla.so looks for libxla_extension.so

exla/c_src/exla/exla.cc

@@ -3,18 +3,18 @@
 #include "exla_nif_util.h"
 #include "exla_client.h"
 #include "exla_log_sink.h"
-#include "tensorflow/compiler/xla/service/platform_util.h"
-#include "tensorflow/compiler/xla/shape_util.h"
-#include "tensorflow/compiler/xla/client/xla_builder.h"
-#include "tensorflow/compiler/xla/client/xla_computation.h"
-#include "tensorflow/compiler/xla/client/client.h"
-#include "tensorflow/compiler/xla/client/lib/math.h"
-#include "tensorflow/compiler/xla/client/lib/lu_decomposition.h"
-#include "tensorflow/compiler/xla/client/lib/qr.h"
-#include "tensorflow/compiler/xla/client/lib/self_adjoint_eig.h"
-#include "tensorflow/compiler/xla/client/lib/svd.h"
-#include "tensorflow/compiler/xla/client/lib/sorting.h"
-#include "tensorflow/compiler/xla/primitive_util.h"
+#include "xla/service/platform_util.h"
+#include "xla/shape_util.h"
+#include "xla/client/xla_builder.h"
+#include "xla/client/xla_computation.h"
+#include "xla/client/client.h"
+#include "xla/client/lib/math.h"
+#include "xla/client/lib/lu_decomposition.h"
+#include "xla/client/lib/qr.h"
+#include "xla/client/lib/self_adjoint_eig.h"
+#include "xla/client/lib/svd.h"
+#include "xla/client/lib/sorting.h"
+#include "xla/primitive_util.h"
 
 // All of these are created with calls to `new` and subsequently
 // passed to the VM as pointers-to-pointers so we balance it out
@@ -2064,7 +2064,7 @@ ERL_NIF_TERM transfer_to_infeed(ErlNifEnv* env, int argc, const ERL_NIF_TERM arg
  xla::Status transfer_status = (*client)->TransferToInfeed(env, terms[0], *shape, device_id);
 
  if(!transfer_status.ok()) {
- return exla::nif::error(env, transfer_status.error_message().c_str());
+ return exla::nif::error(env, transfer_status.message().data());
  }
 
  data = tail;
@@ -2107,7 +2107,7 @@ ERL_NIF_TERM transfer_from_outfeed(ErlNifEnv* env, int argc, const ERL_NIF_TERM
 
  if (!statusor.ok()) {
  enif_clear_env(penv);
- return exla::nif::error(env, statusor.status().error_message().c_str());
+ return exla::nif::error(env, statusor.status().message().data());
  }
 
  ERL_NIF_TERM msg = std::move(statusor.value());
@@ -2326,11 +2326,11 @@ ERL_NIF_TERM start_log_sink(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
  exla::ExlaLogSink* sink = new exla::ExlaLogSink(logger_pid);
 
  // NO_DEFAULT_LOGGER doesn't behave right
- for (auto *log_sink : tensorflow::TFGetLogSinks()) {
- tensorflow::TFRemoveLogSink(log_sink);
+ for (auto *log_sink : tsl::TFGetLogSinks()) {
+ tsl::TFRemoveLogSink(log_sink);
  }
 
- tensorflow::TFAddLogSink(sink);
+ tsl::TFAddLogSink(sink);
 
  return exla::nif::ok(env);
 }

exla/c_src/exla/exla_client.cc

@@ -1,10 +1,10 @@
 #include "exla_client.h"
 #include "exla_nif_util.h"
-#include "tensorflow/compiler/xla/layout_util.h"
-#include "tensorflow/compiler/xla/pjrt/gpu/gpu_helpers.h"
-#include "tensorflow/compiler/xla/pjrt/tfrt_cpu_pjrt_client.h"
-#include "tensorflow/compiler/xla/pjrt/gpu/se_gpu_pjrt_client.h"
-#include "tensorflow/compiler/xla/pjrt/tpu_client.h"
+#include "xla/layout_util.h"
+#include "xla/pjrt/gpu/gpu_helpers.h"
+#include "xla/pjrt/tfrt_cpu_pjrt_client.h"
+#include "xla/pjrt/gpu/se_gpu_pjrt_client.h"
+#include "xla/pjrt/tpu_client.h"
 
 namespace exla {
 
@@ -447,7 +447,7 @@ xla::StatusOr<ExlaClient*> GetGpuClient(double memory_fraction,
  };
 
  EXLA_ASSIGN_OR_RETURN(std::unique_ptr<xla::PjRtClient> client,
- xla::GetStreamExecutorGpuClient(false, allocator_config, nullptr, 0));
+ xla::GetStreamExecutorGpuClient(false, allocator_config, 0));
 
  return new ExlaClient(std::move(client));
 }

exla/c_src/exla/exla_client.h

@@ -6,10 +6,10 @@
 #include <utility>
 
 #include "exla_nif_util.h"
-#include "tensorflow/core/platform/types.h"
-#include "tensorflow/core/platform/status.h"
-#include "tensorflow/compiler/xla/pjrt/gpu/gpu_helpers.h"
-#include "tensorflow/compiler/xla/pjrt/pjrt_client.h"
+#include "tsl/platform/types.h"
+#include "tsl/platform/status.h"
+#include "xla/pjrt/gpu/gpu_helpers.h"
+#include "xla/pjrt/pjrt_client.h"
 
 // The implementations in this module are designed after implementations
 // in the XLA runtime, PjRt. Deviations are made where it makes sense

exla/c_src/exla/exla_log_sink.h

@@ -4,15 +4,15 @@
 #include <string>
 
 #include "exla_nif_util.h"
-#include "tensorflow/core/platform/logging.h"
+#include "tsl/platform/logging.h"
 #include "absl/base/log_severity.h"
 
 namespace exla {
 
 // Redirects calls to logging to the Elixir Logger. `sink_pid`
 // is the PID for a GenServer in Elixir which receives messages
 // with logging information on every call to `LOG(severity)`.
-class ExlaLogSink : public tensorflow::TFLogSink {
+class ExlaLogSink : public tsl::TFLogSink {
  public:
  explicit ExlaLogSink(ErlNifPid sink_pid) : sink_pid_(sink_pid) {
  // Logger Env
@@ -45,7 +45,7 @@ class ExlaLogSink : public tensorflow::TFLogSink {
  return enif_make_tuple4(env_, status, msg, file, line_term);
  }
 
- void Send(const tensorflow::TFLogEntry& entry) {
+ void Send(const tsl::TFLogEntry& entry) {
  ERL_NIF_TERM msg;
  std::string msg_str = entry.ToString();
  std::string fname = entry.FName();

exla/c_src/exla/exla_nif_util.cc

@@ -1,6 +1,6 @@
 #include "exla_nif_util.h"
-#include "tensorflow/compiler/xla/shape_util.h"
-#include "tensorflow/compiler/xla/primitive_util.h"
+#include "xla/shape_util.h"
+#include "xla/primitive_util.h"
 
 namespace exla {
 namespace nif {

exla/c_src/exla/exla_nif_util.h

@@ -9,9 +9,9 @@
 #include <map>
 
 #include "erl_nif.h"
-#include "tensorflow/compiler/xla/xla_data.pb.h"
-#include "tensorflow/compiler/xla/types.h"
-#include "tensorflow/compiler/xla/shape.h"
+#include "xla/xla_data.pb.h"
+#include "xla/types.h"
+#include "xla/shape.h"
 
 #if !defined(__GNUC__) && (defined(__WIN32__) || defined(_WIN32) || defined(_WIN32_))
  typedef unsigned __int64 nif_uint64_t;
@@ -346,7 +346,7 @@ ERL_NIF_TERM make_shape_info(ErlNifEnv* env, xla::Shape shape);
 #define EXLA_EFFECT_OR_RETURN_NIF_IMPL(status, rexpr, env) \
  auto status = (rexpr); \
  if (!status.ok()) { \
- return exla::nif::error(env, status.error_message().c_str()); \
+ return exla::nif::error(env, status.message().data()); \
  }
 
 // Macro to be used to consume Status from within a NIF.
@@ -372,7 +372,7 @@ ERL_NIF_TERM make_shape_info(ErlNifEnv* env, xla::Shape shape);
 #define EXLA_ASSIGN_OR_RETURN_NIF_IMPL(statusor, lhs, rexpr, env) \
  auto statusor = (rexpr); \
  if (!statusor.ok()) { \
- return exla::nif::error(env, statusor.status().error_message().c_str()); \
+ return exla::nif::error(env, statusor.status().message().data()); \
  } \
  lhs = std::move(statusor.value());
 

exla/mix.exs

@@ -62,7 +62,9 @@ defmodule EXLA.MixProject do
  # {:nx, "~> 0.5.1"},
  {:nx, path: "../nx"},
  {:telemetry, "~> 0.4.0 or ~> 1.0"},
- {:xla, "~> 0.4.4", runtime: false},
+ # {:xla, "~> 0.4.4", runtime: false},
+ # {:xla, path: "~/git/xla", runtime: false},
+ {:xla, github: "elixir-nx/xla", runtime: false},
  {:elixir_make, "~> 0.6", runtime: false},
  {:benchee, "~> 1.0", only: :dev},
  {:ex_doc, "~> 0.29.0", only: :docs}

exla/mix.lock

@@ -3,7 +3,7 @@
  "complex": {:hex, :complex, "0.5.0", "af2d2331ff6170b61bb738695e481b27a66780e18763e066ee2cd863d0b1dd92", [:mix], [], "hexpm", "2683bd3c184466cfb94fad74cbfddfaa94b860e27ad4ca1bffe3bff169d91ef1"},
  "deep_merge": {:hex, :deep_merge, "1.0.0", "b4aa1a0d1acac393bdf38b2291af38cb1d4a52806cf7a4906f718e1feb5ee961", [:mix], [], "hexpm", "ce708e5f094b9cd4e8f2be4f00d2f4250c4095be93f8cd6d018c753894885430"},
  "earmark_parser": {:hex, :earmark_parser, "1.4.31", "a93921cdc6b9b869f519213d5bc79d9e218ba768d7270d46fdcf1c01bacff9e2", [:mix], [], "hexpm", "317d367ee0335ef037a87e46c91a2269fef6306413f731e8ec11fc45a7efd059"},
- "elixir_make": {:hex, :elixir_make, "0.7.4", "5439110c964ffdd8212ca919b5b8beac423085a77ad33d5e394abe812c2d2d75", [:mix], [{:castore, "~> 0.1", [hex: :castore, repo: "hexpm", optional: true]}], "hexpm", "70c33052f7b00c813fd66d15a3cf1f7d1e122860c572ec81b8181b1276074157"},
+ "elixir_make": {:hex, :elixir_make, "0.7.7", "7128c60c2476019ed978210c245badf08b03dbec4f24d05790ef791da11aa17c", [:mix], [{:castore, "~> 0.1 or ~> 1.0", [hex: :castore, repo: "hexpm", optional: true]}], "hexpm", "5bc19fff950fad52bbe5f211b12db9ec82c6b34a9647da0c2224b8b8464c7e6c"},
  "ex_doc": {:hex, :ex_doc, "0.29.3", "f07444bcafb302db86e4f02d8bbcd82f2e881a0dcf4f3e4740e4b8128b9353f7", [:mix], [{:earmark_parser, "~> 1.4.31", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_elixir, "~> 0.14", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1", [hex: :makeup_erlang, repo: "hexpm", optional: false]}], "hexpm", "3dc6787d7b08801ec3b51e9bd26be5e8826fbf1a17e92d1ebc252e1a1c75bfe1"},
  "makeup": {:hex, :makeup, "1.1.0", "6b67c8bc2882a6b6a445859952a602afc1a41c2e08379ca057c0f525366fc3ca", [:mix], [{:nimble_parsec, "~> 1.2.2 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "0a45ed501f4a8897f580eabf99a2e5234ea3e75a4373c8a52824f6e873be57a6"},
  "makeup_elixir": {:hex, :makeup_elixir, "0.16.0", "f8c570a0d33f8039513fbccaf7108c5d750f47d8defd44088371191b76492b0b", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}, {:nimble_parsec, "~> 1.2.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "28b2cbdc13960a46ae9a8858c4bebdec3c9a6d7b4b9e7f4ed1502f8159f338e7"},
@@ -12,5 +12,5 @@
  "nx": {:hex, :nx, "0.5.1", "118134b8c97c2a8f86c87aa8434994c1cbbe139a306b89cca04e08dd46228067", [:mix], [{:complex, "~> 0.5", [hex: :complex, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "ceb8fbbe19b3c4252a7188d8b0e059fac9da0f4a4f3bb770fc665fdd0b29f0c5"},
  "statistex": {:hex, :statistex, "1.0.0", "f3dc93f3c0c6c92e5f291704cf62b99b553253d7969e9a5fa713e5481cd858a5", [:mix], [], "hexpm", "ff9d8bee7035028ab4742ff52fc80a2aa35cece833cf5319009b52f1b5a86c27"},
  "telemetry": {:hex, :telemetry, "1.2.1", "68fdfe8d8f05a8428483a97d7aab2f268aaff24b49e0f599faa091f1d4e7f61c", [:rebar3], [], "hexpm", "dad9ce9d8effc621708f99eac538ef1cbe05d6a874dd741de2e689c47feafed5"},
- "xla": {:hex, :xla, "0.4.4", "c3a8ed1f579bda949df505e49ff65415c8281d991fbd6ae1d8f3c5d0fd155f54", [:make, :mix], [{:elixir_make, "~> 0.4", [hex: :elixir_make, repo: "hexpm", optional: false]}], "hexpm", "484f3f9011db3c9f1ff1e98eecefd382f3882a07ada540fd58803db1d2dab671"},
+ "xla": {:git, "https://github.com/elixir-nx/xla.git", "a864469e7a5ecec4569cd260b50347c1a274deaf", []},
 }