fix: update iree

polvalente · polvalente · commit 413d754d7a84 · 2024-08-23T03:28:52.000-03:00
diff --git a/.github/workflows/embedded_devices.yml b/.github/workflows/embedded_devices.yml
@@ -16,7 +16,7 @@ jobs:
       MIX_ENV: prod
       NX_IREE_PREFER_PRECOMPILED: false
       NX_IREE_SOURCE_DIR: ./build-cache/iree
-      IREE_GIT_REV: candidate-20240604.914
+      IREE_GIT_REV: candidate-20240822.993
     strategy:
       fail-fast: true
       matrix:
diff --git a/.github/workflows/precompiled_nif.yml b/.github/workflows/precompiled_nif.yml
@@ -16,7 +16,7 @@ jobs:
       MIX_ENV: prod
       NX_IREE_PREFER_PRECOMPILED: false
       NX_IREE_SOURCE_DIR: ./build-cache/iree
-      IREE_GIT_REV: candidate-20240604.914
+      IREE_GIT_REV: candidate-20240822.993
       BUILD_IREE_RUNTIME: false
     strategy:
       fail-fast: false
@@ -83,7 +83,7 @@ jobs:
       MIX_ENV: prod
       NX_IREE_PREFER_PRECOMPILED: false
       NX_IREE_SOURCE_DIR: ./build-cache/iree
-      IREE_GIT_REV: candidate-20240604.914
+      IREE_GIT_REV: candidate-20240822.993
       ImageOS: ubuntu22
       LANG: en_US.UTF-8
       LANGUAGE: en_US:en
diff --git a/compiler.exs b/compiler.exs
@@ -18,6 +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)
 arg0 = Nx.tensor([1.0, 2.0, 3.0, 4.0])
 arg1 = Nx.tensor([1, -1, 1, -1])
 f.(arg0, arg1) |> dbg()
diff --git a/lib/nx_iree.ex b/lib/nx_iree.ex
@@ -69,7 +69,11 @@ defmodule NxIREE do
 
     [driver_name, _] = String.split(device, "://", parts: 2)
 
-    input_refs = Enum.map(inputs, &NxIREE.VM.allocate_buffer(&1, device_ref))
+    input_refs =
+      Enum.map(inputs, fn
+        %Nx.Tensor{data: %NxIREE.Tensor{ref: ref}} -> ref
+        t -> NxIREE.VM.allocate_buffer(t, device_ref)
+      end)
 
     instance_ref = NxIREE.VM.get_instance()
 
diff --git a/liveview_native/live_nx_iree/lib/live_nx_iree_web/live/home_live/home_live.ex b/liveview_native/live_nx_iree/lib/live_nx_iree_web/live/home_live/home_live.ex
@@ -65,10 +65,10 @@ defmodule LiveNxIREEWeb.HomeLive do
     {backend_flag, runtime_device} =
       case target_device do
         :metal ->
-          {"--iree-hal-target-backends=metal-spirv", "metal://default"}
+          {"--iree-hal-target-backends=metal-spirv", "metal://"}
 
         :cpu ->
-          {"--iree-hal-target-backends=llvm-cpu", "local-sync://default"}
+          {"--iree-hal-target-backends=llvm-cpu", "local-sync://"}
       end
 
     compiler_flags = [
@@ -104,7 +104,7 @@ defmodule LiveNxIREEWeb.HomeLive do
 
       {:ok, serialized} = NxIREE.Native.serialize_tensor(tensor.data.ref)
 
-      serialized
+      Base.encode64(serialized)
     end)
   end
 
diff --git a/liveview_native/live_nx_iree/lib/live_nx_iree_web/live/swiftui/home_live.swiftui.neex b/liveview_native/live_nx_iree/lib/live_nx_iree_web/live/swiftui/home_live.swiftui.neex
@@ -1 +1,7 @@
-<NxFunction on-execution="nx-executed" signature={@function_signature} bytecode={@bytecode} device={@device} inputs={@inputs} num-outputs={@num_outputs} />
+<NxFunction
+  on-execution="nx-executed"
+  signature={@function_signature}
+  bytecode={@bytecode}
+  device={@device}
+  inputs={@inputs}
+  num-outputs={@num_outputs} />
diff --git a/liveview_native/live_nx_iree/native/swiftui/LiveNxIREE.xcodeproj/project.pbxproj b/liveview_native/live_nx_iree/native/swiftui/LiveNxIREE.xcodeproj/project.pbxproj
diff --git a/liveview_native/live_nx_iree/native/swiftui/LiveNxIREE.xcodeproj/project.xcworkspace/xcuserdata/paulo.valente.xcuserdatad/UserInterfaceState.xcuserstate b/liveview_native/live_nx_iree/native/swiftui/LiveNxIREE.xcodeproj/project.xcworkspace/xcuserdata/paulo.valente.xcuserdatad/UserInterfaceState.xcuserstate
diff --git a/liveview_native/live_nx_iree/native/swiftui/LiveNxIREE/LiveNxIREE.swift b/liveview_native/live_nx_iree/native/swiftui/LiveNxIREE/LiveNxIREE.swift
@@ -7,26 +7,6 @@ import SwiftUI
 
 @main
 struct LiveNxIREE: App {
-    init() {
-        // Allocate memory for the pointers
-        let vmInstance = UnsafeMutablePointer<iree_vm_instance_t>.allocate(capacity: 1)
-        let driverRegistry = UnsafeMutablePointer<iree_hal_driver_registry_t>.allocate(capacity: 1)
-        let errorMessage = UnsafeMutablePointer<CChar>.allocate(capacity: 256)
-                
-        // Call the initialization function
-        let result = nx_iree_initialize(vmInstance, driverRegistry, errorMessage)
-            
-        if result != 0 {
-            // Handle the error
-            let errorString = String(cString: errorMessage)
-            print("Error initializing nx_iree: \(errorString)")
-        } else {
-            globalVmInstance = vmInstance
-            globalDriverRegistry = driverRegistry
-            print("nx_iree initialized successfully.")
-        }
-    }
-    
     var body: some Scene {
         WindowGroup {
             ContentView()
diff --git a/liveview_native/live_nx_iree/native/swiftui/LiveNxIREE/NxAddon.swift b/liveview_native/live_nx_iree/native/swiftui/LiveNxIREE/NxAddon.swift
@@ -9,34 +9,24 @@ import LiveViewNative
 import SwiftUI
 
 class iree_vm_instance_t {}
-class iree_hal_driver_registry_t {}
 class iree_hal_device_t {}
 
-var globalVmInstance: UnsafeMutablePointer<iree_vm_instance_t>?
-var globalDriverRegistry: UnsafeMutablePointer<iree_hal_driver_registry_t>?
-
-@_silgen_name("nx_iree_initialize")
-func nx_iree_initialize(
-    _ vm_instance: UnsafeMutablePointer<iree_vm_instance_t>,
-    _ driver_registry: UnsafeMutablePointer<iree_hal_driver_registry_t>,
-    _ error_message: UnsafeMutablePointer<CChar>) -> Int
+@_silgen_name("nx_iree_create_instance")
+func nx_iree_create_instance() -> UnsafePointer<iree_vm_instance_t>?
 
 @_silgen_name("nx_iree_create_device")
-func nx_iree_create_device(
-    _ driver_registry: UnsafeMutablePointer<iree_hal_driver_registry_t>,
-    _ name: UnsafePointer<CChar>) -> UnsafeMutablePointer<iree_hal_device_t>
+func nx_iree_create_device(_ name: UnsafePointer<CChar>) -> UnsafePointer<iree_hal_device_t>?
 
 @_silgen_name("nx_iree_call")
 func nx_iree_call(
-    _ vm_instance: UnsafeMutablePointer<iree_vm_instance_t>,
-    _ device: UnsafeMutablePointer<iree_hal_device_t>,
+    _ vm_instance: UnsafePointer<iree_vm_instance_t>,
+    _ device: UnsafePointer<iree_hal_device_t>,
     _ bytecode_size: UInt64,
     _ bytecode: UnsafePointer<CUnsignedChar>,
     _ num_inputs: UInt64,
     _ serialized_inputs: UnsafePointer<UnsafePointer<CChar>>,
     _ num_outputs: UInt64,
-    _ serialized_outputs: UnsafePointer<UnsafePointer<CChar>>,
-    _ error_message: UnsafeMutablePointer<CChar>) -> Int
+    _ error_message: UnsafeMutablePointer<CChar>) -> UnsafePointer<UnsafePointer<CChar>>?
 
 
 
diff --git a/liveview_native/live_nx_iree/native/swiftui/LiveNxIREE/NxAddon/NxFunctionView.swift b/liveview_native/live_nx_iree/native/swiftui/LiveNxIREE/NxAddon/NxFunctionView.swift
@@ -18,6 +18,13 @@ struct NxFunctionView<Root: RootRegistry>: View {
     @LiveAttribute("num-outputs") private var numOutputs: Int? = nil
     @Event("on-execution", type: "change") private var change
     
+    @LiveElementIgnored
+    private var vmInstance: UnsafePointer<iree_vm_instance_t>? = nil
+    
+    init() {
+        vmInstance = nx_iree_create_instance()
+    }
+    
     var body: some View {
         VStack {
             if signature != nil {
@@ -115,43 +122,48 @@ struct NxFunctionView<Root: RootRegistry>: View {
     private func run() {
         if bytecode != nil,
            deviceURI != nil,
-           globalVmInstance != nil,
-           globalDriverRegistry != nil,
+           vmInstance != nil,
            serializedInputs != nil,
            let (bytecodeSize, bytecodePointer) = convertBase64StringToBytecode(bytecode!),
            let inputs = convertToCStringArray(from: serializedInputs!) {
             let deviceURIcstr = strdup(deviceURI!)
-            let device = nx_iree_create_device(globalDriverRegistry!, UnsafePointer(deviceURIcstr)!)
+            let device = nx_iree_create_device(UnsafePointer(deviceURIcstr)!)
             deviceURIcstr?.deallocate()
             
-            let serializedOutputs: UnsafePointer<UnsafePointer<CChar>>? = nil
             let errorMessage = UnsafeMutablePointer<CChar>.allocate(capacity: 256)
             
             print("Executing function \(signature ?? "None") on device: \(deviceURI ?? "None")")
 
-            let result = nx_iree_call(
-                globalVmInstance!,
-                device,
+            let serializedOutputs = nx_iree_call(
+                vmInstance!,
+                device!,
                 bytecodeSize,
                 bytecodePointer!,
                 UInt64(serializedInputs!.count),
                 inputs,
                 UInt64(numOutputs!),
-                serializedOutputs!,
                 errorMessage)
             
-            if result != 0 {
-                print(errorMessage)
+            if serializedOutputs == nil {
+                let errorString = String(cString: errorMessage)
+                print("Failed execution with error: \(errorString)")
                 return
             }
             
             for i in 0..<serializedInputs!.count {
                 inputs[i].deallocate()
-           }
+            }
             inputs.deallocate()
             
+            print("Finished executing. Preparing to send results back.")
+            print("outputs: \(serializedOutputs)")
+            print("numOutputs: \(numOutputs)")
+            
             change(value: base64EncodedStrings(from: serializedOutputs!, count: numOutputs!))
         } else {
+            print("vm: \(vmInstance)")
+            print("deviceURI: \(deviceURI)")
+            print("bytecode: \(bytecode)")
             print("IREE components are not initialized.")
         }
    }
diff --git a/liveview_native/live_nx_iree/native/swiftui/LiveNxIREE/c_src/nx_iree.cpp b/liveview_native/live_nx_iree/native/swiftui/LiveNxIREE/c_src/nx_iree.cpp
@@ -28,11 +28,18 @@ int nx_iree_initialize(iree_vm_instance_t* vm_instance, iree_hal_driver_registry
     return 0;
 }
 
-iree_hal_device_t* nx_iree_create_device(iree_hal_driver_registry_t* registry, char* device_uri) {
-    return create_device(registry, std::string(device_uri));
+iree_hal_device_t* nx_iree_create_device(char* device_uri) {
+    iree_hal_driver_registry_t* driver_registry = get_driver_registry();
+    register_all_drivers(driver_registry);
+    
+    return create_device(driver_registry, std::string(device_uri));
+}
+
+iree_vm_instance_t* nx_iree_create_instance() {
+    return create_instance();
 }
 
-int nx_iree_call(iree_vm_instance_t* vm_instance, iree_hal_device_t* device, uint64_t bytecode_size, unsigned char* bytecode, uint64_t num_inputs, char** serialized_inputs, uint64_t num_outputs, char** serialized_outputs, char* error_message) {
+char** nx_iree_call(iree_vm_instance_t* vm_instance, iree_hal_device_t* device, uint64_t bytecode_size, unsigned char* bytecode, uint64_t num_inputs, char** serialized_inputs, uint64_t num_outputs, char* error_message) {
     std::vector<iree::runtime::IREETensor*> inputs;
     
     for (size_t i = 0; i < num_inputs; i++) {
@@ -43,15 +50,12 @@ int nx_iree_call(iree_vm_instance_t* vm_instance, iree_hal_device_t* device, uin
     auto [status, optional_result] = call(vm_instance, device, "not_cuda", bytecode, static_cast<size_t>(bytecode_size), inputs);
         
     if (!is_ok(status)) {
-        if (error_message) {
-            std::string msg = get_status_message(status);
-            strncpy(error_message, msg.c_str(), msg.length());
-        }
-        
-        return 1;
+        std::string msg = get_status_message(status);
+        strncpy(error_message, msg.c_str(), msg.length());
+        return nullptr;
     }
     
-    serialized_outputs = new char*[num_outputs];
+    auto serialized_outputs = new char*[num_outputs];
     
     auto result = optional_result.value();
     for (size_t i = 0; i < num_outputs; i++) {
@@ -60,6 +64,6 @@ int nx_iree_call(iree_vm_instance_t* vm_instance, iree_hal_device_t* device, uin
         memcpy(serialized_outputs[i], serialized->data(), serialized->size());
     }
     
-    return 0;
+    return serialized_outputs;
 }
 
diff --git a/liveview_native/live_nx_iree/native/swiftui/LiveNxIREE/c_src/nx_iree.h b/liveview_native/live_nx_iree/native/swiftui/LiveNxIREE/c_src/nx_iree.h
@@ -15,9 +15,9 @@
 extern "C" {
 #endif
 
-int nx_iree_initialize(iree_vm_instance_t* vm_instance, iree_hal_driver_registry_t* driver_registry, char* error_message);
-iree_hal_device_t* nx_iree_create_device(iree_hal_driver_registry_t* registry, char* device_uri);
-int nx_iree_call(iree_vm_instance_t* vm_instance, iree_hal_device_t* device, uint64_t bytecode_size, unsigned char* bytecode, uint64_t num_inputs, char** serialized_inputs, uint64_t num_outputs, char** serialized_outputs, char* error_message);
+iree_vm_instance_t* nx_iree_create_instance();
+iree_hal_device_t* nx_iree_create_device(char* device_uri);
+char** nx_iree_call(iree_vm_instance_t* vm_instance, iree_hal_device_t* device, uint64_t bytecode_size, unsigned char* bytecode, uint64_t num_inputs, char** serialized_inputs, uint64_t num_outputs, char* error_message);
 
 
 #ifdef __cplusplus
diff --git a/liveview_native/live_nx_iree/native/swiftui/LiveNxIREE/nx_iree/include/nx_iree/runtime.h b/liveview_native/live_nx_iree/native/swiftui/LiveNxIREE/nx_iree/include/nx_iree/runtime.h
@@ -74,7 +74,7 @@ class IREETensor {
 
 iree_vm_instance_t* create_instance();
 iree_hal_driver_registry_t* get_driver_registry();
-iree_hal_device_t* create_device(iree_hal_driver_registry_t *registry, const std::string& device_uri);
+iree_hal_device_t* create_device(iree_hal_driver_registry_t* registry, const std::string& device_uri);
 
 std::pair<iree_status_t, std::optional<std::vector<iree::runtime::IREETensor*>>>
 call(iree_vm_instance_t* i, iree_hal_device_t*, std::string, unsigned char*, size_t, std::vector<iree::runtime::IREETensor*>);
diff --git a/mix.exs b/mix.exs
@@ -71,8 +71,8 @@ defmodule NxIREE.MixProject do
   end
 
   defp nx_iree_config() do
-    version = System.get_env("NX_IREE_VERSION", "20240818.989")
-    tag = System.get_env("NX_IREE_GIT_REV", "candidate-20240818.989")
+    version = System.get_env("NX_IREE_VERSION", "20240822.993")
+    tag = System.get_env("NX_IREE_GIT_REV", "candidate-20240822.993")
 
     env_dir = System.get_env("NX_IREE_COMPILER_DIR")
 
