Merge pull request #17 from LaurentMazare/0.7.0

Update for 0.7.0.

.github/workflows/rust-ci.yml

@@ -18,7 +18,7 @@ jobs:
           toolchain: ${{ matrix.rust }}
           override: true
       - run: |
-          wget -nv https://github.com/elixir-nx/xla/releases/download/v0.5.1/xla_extension-x86_64-linux-gnu-cpu.tar.gz
+          wget -nv https://github.com/elixir-nx/xla/releases/download/v0.7.0/xla_extension-x86_64-linux-gnu-cpu.tar.gz
           tar -xzvf xla_extension-x86_64-linux-gnu-cpu.tar.gz
       - uses: actions-rs/cargo@v1
         with:
@@ -39,7 +39,7 @@ jobs:
           toolchain: ${{ matrix.rust }}
           override: true
       - run: |
-          wget -nv https://github.com/elixir-nx/xla/releases/download/v0.5.1/xla_extension-x86_64-linux-gnu-cpu.tar.gz
+          wget -nv https://github.com/elixir-nx/xla/releases/download/v0.7.0/xla_extension-x86_64-linux-gnu-cpu.tar.gz
           tar -xzvf xla_extension-x86_64-linux-gnu-cpu.tar.gz
       - uses: actions-rs/cargo@v1
         with:
@@ -73,7 +73,7 @@ jobs:
           override: true
       - run: rustup component add clippy
       - run: |
-          wget -nv https://github.com/elixir-nx/xla/releases/download/v0.5.1/xla_extension-x86_64-linux-gnu-cpu.tar.gz
+          wget -nv https://github.com/elixir-nx/xla/releases/download/v0.7.0/xla_extension-x86_64-linux-gnu-cpu.tar.gz
           tar -xzvf xla_extension-x86_64-linux-gnu-cpu.tar.gz
       - uses: actions-rs/cargo@v1
         with:

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "xla"
-version = "0.1.6"
+version = "0.2.0"
 authors = ["laurent <[email protected]>"]
 edition = "2021"
 description = "Bindings for the XLA C++ library."

README.md

@@ -2,14 +2,14 @@
 Experimentation using the xla compiler from rust
 
 Pre-compiled binaries for the xla library can be downloaded from the
-[elixir-nx/xla repo](https://github.com/elixir-nx/xla/releases/tag/v0.6.0).
+[elixir-nx/xla repo](https://github.com/elixir-nx/xla/releases/tag/v0.7.0).
 These should be extracted at the root of this repository, resulting
 in a `xla_extension` subdirectory being created, the currently supported version
-is 0.6.0.
+is 0.7.0.
 
 For a linux platform, this can be done via:
 ```bash
-wget https://github.com/elixir-nx/xla/releases/download/v0.6.0/xla_extension-x86_64-linux-gnu-cpu.tar.gz
+wget https://github.com/elixir-nx/xla/releases/download/v0.7.0/xla_extension-x86_64-linux-gnu-cpu.tar.gz
 tar -xzvf xla_extension-x86_64-linux-gnu-cpu.tar.gz
 ```
 

build.rs

@@ -6,6 +6,7 @@ use std::path::{Path, PathBuf};
 #[derive(Clone, Copy, Eq, PartialEq)]
 enum OS {
     Linux,
+    #[allow(clippy::enum_variant_names)]
     MacOS,
     Windows,
 }
@@ -80,7 +81,7 @@ fn main() {
     // "DSO missing from command line" error
     // undefined reference to symbol '_ZStlsIcSt11char_traitsIcESaIcEERSt13basic_ostreamIT_T0_ES7_RKNSt7__cxx1112basic_stringIS4_S5_T1_EE@@GLIBCXX_3.4.21'
     if os == OS::Linux {
-        println!("cargo:rustc-link-arg=-Wl,--copy-dt-needed-entries");
+        // println!("cargo:rustc-link-arg=-Wl,--copy-dt-needed-entries");
         println!("cargo:rustc-link-arg=-Wl,-lstdc++");
     }
     println!("cargo:rustc-link-search=native={}", xla_dir.join("lib").display());

xla_rs/xla_rs.cc

@@ -42,8 +42,9 @@ status pjrt_gpu_client_create(pjrt_client *output, double memory_fraction,
                               bool preallocate) {
   xla::GpuAllocatorConfig allocator = {.memory_fraction = memory_fraction,
                                        .preallocate = preallocate};
-  ASSIGN_OR_RETURN_STATUS(
-      client, xla::GetStreamExecutorGpuClient(false, allocator, 0, 0));
+  xla::GpuClientOptions options;
+  options.allocator_config = allocator;
+  ASSIGN_OR_RETURN_STATUS(client, xla::GetStreamExecutorGpuClient(options));
   *output = new std::shared_ptr(std::move(client));
   return nullptr;
 }