feat: add full Nx backend (#11)

* feat: add full backend

* feat: run nx test suite

* feat: set default device

* fix: recover from partical iree clones

* Update cmake/src/runtime.cc

Author: polvalente

Makefile

@@ -23,7 +23,8 @@ endif
 .PHONY: clone_iree
 clone_iree: $(NX_IREE_SOURCE_DIR)
 
-$(NX_IREE_SOURCE_DIR):
+.PHONY: iree_source_dir
+iree_source_dir:
 	./scripts/clone_iree.sh $(IREE_GIT_REV) $(NX_IREE_SOURCE_DIR)
 
 IREE_CMAKE_BUILD_DIR ?= $(abspath iree-runtime/iree-build)
@@ -106,7 +107,7 @@ install_runtime: iree_host $(IREE_INSTALL_DIR)
 
 CMAKE_SOURCES = $(abspath cmake/src/runtime.cc) $(abspath cmake/src/runtime.h)
 
-$(IREE_INSTALL_DIR): $(NX_IREE_SOURCE_DIR) $(CMAKE_SOURCES)
+$(IREE_INSTALL_DIR): iree_source_dir $(CMAKE_SOURCES)
 	cmake -G Ninja -B $(IREE_CMAKE_BUILD_DIR) \
 		-DCMAKE_BUILD_TYPE=$(IREE_CMAKE_CONFIG)\
 		-DIREE_BUILD_COMPILER=OFF\

c_src/nx_iree.cc

@@ -13,6 +13,10 @@ ERL_NIF_TERM error(ErlNifEnv* env, const char* error) {
   return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_string(env, error, ERL_NIF_LATIN1));
 }
 
+ERL_NIF_TERM ok(ErlNifEnv* env) {
+  return enif_make_atom(env, "ok");
+}
+
 ERL_NIF_TERM ok(ErlNifEnv* env, ERL_NIF_TERM term) {
   return enif_make_tuple2(env, enif_make_atom(env, "ok"), term);
 }
@@ -233,7 +237,8 @@ DECLARE_NIF(list_devices) {
     auto ref_term = make<iree_hal_device_t*>(env, device->ref);
     auto driver_name_term = enif_make_string(env, device->driver_name.c_str(), ERL_NIF_LATIN1);
     auto uri_term = enif_make_string(env, device->uri.c_str(), ERL_NIF_LATIN1);
-    auto tuple = enif_make_tuple3(env, ref_term, driver_name_term, uri_term);
+    auto id_term = enif_make_uint64(env, device->id);
+    auto tuple = enif_make_tuple4(env, ref_term, driver_name_term, uri_term, id_term);
     device_terms.push_back(tuple);
   }
 
@@ -315,13 +320,13 @@ iree_hal_element_type_t nx_type_to_iree_type(std::string type) {
     return type_enum::IREE_HAL_ELEMENT_TYPE_INT_32;
   } else if (type == "i64") {
     return type_enum::IREE_HAL_ELEMENT_TYPE_INT_64;
-  } else if (type == "u8") {
+  } else if (type == "ui8") {
     return type_enum::IREE_HAL_ELEMENT_TYPE_UINT_8;
-  } else if (type == "u16") {
+  } else if (type == "ui16") {
     return type_enum::IREE_HAL_ELEMENT_TYPE_UINT_16;
-  } else if (type == "u32") {
+  } else if (type == "ui32") {
     return type_enum::IREE_HAL_ELEMENT_TYPE_UINT_32;
-  } else if (type == "u64") {
+  } else if (type == "ui64") {
     return type_enum::IREE_HAL_ELEMENT_TYPE_UINT_64;
   } else if (type == "bf16") {
     return type_enum::IREE_HAL_ELEMENT_TYPE_BFLOAT_16;
@@ -355,14 +360,10 @@ DECLARE_NIF(read_buffer_nif) {
     return error(env, "invalid num_bytes");
   }
 
-  std::cout << "num_bytes input: " << num_bytes << std::endl;
-
   if (num_bytes == -1) {
     num_bytes = (*input)->size;
   }
 
-  std::cout << "num_bytes actual: " << num_bytes << std::endl;
-
   ErlNifBinary binary;
 
   if (!enif_alloc_binary(num_bytes, &binary)) {
@@ -418,6 +419,22 @@ DECLARE_NIF(allocate_buffer) {
   return ok(env, make<iree::runtime::IREETensor*>(env, input));
 }
 
+DECLARE_NIF(deallocate_buffer) {
+  if (argc != 1) {
+    return error(env, "invalid number of arguments");
+  }
+
+  iree::runtime::IREETensor** input;
+
+  if (!get<iree::runtime::IREETensor*>(env, argv[0], input)) {
+    return error(env, "invalid input");
+  }
+
+  (*input)->deallocate();
+
+  return ok(env);
+}
+
 DECLARE_NIF(serialize_tensor) {
   if (argc != 1) {
     return error(env, "invalid number of arguments");
@@ -510,6 +527,7 @@ static ErlNifFunc funcs[] = {
     {"list_devices", 1, list_devices},
     {"list_devices", 2, list_devices},
     {"list_drivers", 1, list_drivers},
+    {"deallocate_buffer", 1, deallocate_buffer},
     {"allocate_buffer", 4, allocate_buffer},
     {"serialize_tensor", 1, serialize_tensor},
     {"deserialize_tensor", 1, deserialize_tensor},

cmake/src/runtime.cc

@@ -70,7 +70,24 @@ iree::runtime::IREETensor::IREETensor(char *buffer) {
   this->buffer_view = nullptr;
 }
 
-std::vector<char> *iree::runtime::IREETensor::serialize() {
+iree::runtime::IREETensor::~IREETensor() {
+  this->deallocate();
+}
+
+void iree::runtime::IREETensor::deallocate() {
+  if (data) {
+    std::free(data);
+    data = nullptr;
+  }
+
+  if (buffer_view) {
+    iree_hal_buffer_view_release(buffer_view);
+    buffer_view = nullptr;
+  }
+}
+
+std::vector<char> *
+iree::runtime::IREETensor::serialize() {
   auto buffer = new std::vector<char>();
 
   // Serialize 'type'
@@ -196,16 +213,6 @@ iree_status_t list_devices(iree_hal_driver_registry_t *registry, std::string dri
     return status;
   }
 
-  auto out_device = new iree::runtime::Device(driver_name);
-  status = iree_hal_driver_create_default_device(driver, iree_allocator_system(),
-                                                 &out_device->ref);
-  if (!iree_status_is_ok(status)) {
-    return status;
-  }
-
-  out_device->uri = driver_name + "://default";
-  devices.push_back(out_device);
-
   status = iree_hal_driver_query_available_devices(
       driver, iree_allocator_system(), &device_info_count, &device_infos);
 
@@ -217,7 +224,9 @@ iree_status_t list_devices(iree_hal_driver_registry_t *registry, std::string dri
   for (size_t i = 0; i < device_info_count; i++) {
     auto device = new iree::runtime::Device(driver_name);
     auto info = device_infos[i];
-    device->uri = driver_name + "://" + std::string(info.path.data, info.path.size);
+    std::string device_urn(info.path.data, info.path.size);
+    device->uri = driver_name + "://" + device_urn;
+    device->id = info.device_id;
 
     status = iree_hal_create_device(
         registry,

cmake/src/runtime.h

@@ -29,6 +29,7 @@ class Device {
   std::string uri;
   iree_hal_device_t* ref;
   std::string driver_name;
+  uint64_t id;
 
   Device(std::string driver_name) : driver_name(driver_name) {}
   ~Device();
@@ -48,12 +49,9 @@ class IREETensor {
   IREETensor(void* data, size_t size, std::vector<int64_t> in_dims, iree_hal_element_type_t type);
 
   // Destructor
-  ~IREETensor() {
-    if (data) {
-      std::free(data);
-      data = nullptr;
-    }
-  }
+  ~IREETensor();
+
+  void deallocate();
 
   // Disable copy and move semantics for simplicity
   IREETensor(const IREETensor&) = delete;

compiler.exs

@@ -18,7 +18,7 @@ Nx.Defn.default_options(compiler: NxIREE.Compiler, iree_compiler_flags: flags, i
 
 f = Nx.Defn.compile(fun, args)
 
-Nx.default_backend(NxIREE.Tensor)
+Nx.default_backend(NxIREE.Backend)
 arg0 = Nx.tensor([1.0, 2.0, 3.0, 4.0])
 arg1 = Nx.tensor([1, -1, 1, -1])
 f.(arg0, arg1) |> dbg()

config/config.exs

@@ -0,0 +1,3 @@
+import Config
+
+import_config "#{config_env()}.exs"

config/dev.exs

@@ -0,0 +1,3 @@
+import Config
+
+config :nx_iree, :add_backend_on_inspect, false

config/prod.exs

@@ -20,7 +20,7 @@ if System.get_env("PHX_SERVER") do
   config :live_nx_iree, LiveNxIREEWeb.Endpoint, server: true
 end
 
-config :nx, :default_backend, NxIREE.Tensor
+config :nx, :default_backend, NxIREE.Backend
 
 if config_env() == :prod do
   database_url =