feat: send bytecode and signature to device given handle_info message

Author: polvalente

lib/nx_iree/compiler.ex

@@ -3,6 +3,15 @@ defmodule NxIREE.Compiler do
   Compiler for Nx defn
   """
 
+  def to_bytecode(fun, templates, opts \\ []) do
+    opts = opts |> Keyword.put(:output_mode, :bytecode) |> Keyword.put(:compiler, __MODULE__)
+
+    Nx.Defn.compile(fun, templates, opts)
+  catch
+    {:bytecode, %{bytecode: bytecode, output_container: output_container}} ->
+      {:ok, %{bytecode: bytecode, output_container: output_container}}
+  end
+
   @behaviour Nx.Defn.Compiler
 
   @impl true
@@ -22,9 +31,7 @@ defmodule NxIREE.Compiler do
     bytecode = NxIREE.compile(mlir_module, iree_compiler_flags)
 
     if output_mode == :bytecode do
-      fn _ ->
-        [bytecode]
-      end
+      throw({:bytecode, %{bytecode: bytecode, output_container: output_container}})
     else
       fn [inputs] ->
         {:ok, results} =

liveview_native/live_nx_iree/config/config.exs

@@ -9,7 +9,7 @@ import Config
 
 config :live_nx_iree,
   namespace: LiveNxIREE,
-  ecto_repos: [LiveNxIREE.Repo],
+  ecto_repos: [],
   generators: [timestamp_type: :utc_datetime]
 
 # Configures the endpoint

liveview_native/live_nx_iree/lib/live_nx_iree/application.ex

@@ -9,7 +9,7 @@ defmodule LiveNxIREE.Application do
   def start(_type, _args) do
     children = [
       LiveNxIREEWeb.Telemetry,
-      LiveNxIREE.Repo,
+      # LiveNxIREE.Repo,
       {DNSCluster, query: Application.get_env(:live_nx_iree, :dns_cluster_query) || :ignore},
       {Phoenix.PubSub, name: LiveNxIREE.PubSub},
       # Start the Finch HTTP client for sending emails

liveview_native/live_nx_iree/lib/live_nx_iree_web/live/home_live/home_live.ex

@@ -4,6 +4,19 @@ defmodule LiveNxIREEWeb.HomeLive do
 
   @impl true
   def mount(_params, _session, socket) do
+    # register the liveview with presence.
+    # This will make it so that we can interact with the liveview
+    # from outside of the liveview itself.
+
+    # this can probably be achieved by just changing assigns from within a handle_info call, somehow
+
+    # need to figure out how to push and pull events from LVN. otherwise, we'll fallback
+    # to a long-polling approach from within the liveview itself using a jobqueue via pubsub.
+
+    dbg(self())
+
+    socket = assign(socket, bytecode: nil, function_signature: nil)
+
     {:ok, socket}
   end
 
@@ -17,4 +30,71 @@ defmodule LiveNxIREEWeb.HomeLive do
     |> assign(:page_title, "Listing Contexts")
     |> assign(:home, nil)
   end
+
+  # def handle_info({:nx, :execute, function, input_templates}, socket) do
+  #   socket =
+  #     assign(
+  #       socket,
+  #       :bytecode,
+  #       Base.encode64(inspect({:nx, :execute, function, input_templates}))
+  #     )
+
+  #   {:noreply, socket}
+  # end
+
+  @impl true
+  def handle_info({:nx, :execute, function, input_templates, target_device}, socket) do
+    fun =
+      case function do
+        {m, f, a} ->
+          Function.capture(m, f, a)
+
+        _ ->
+          raise """
+          Expected a tuple of module, function, and arguments, but got: #{inspect(function)}
+          """
+      end
+
+    backend_flag =
+      case target_device do
+        :metal -> "--iree-hal-target-backends=metal-spirv"
+        :cpu -> "--iree-hal-target-backends=llvm-cpu"
+      end
+
+    compiler_flags = [
+      backend_flag,
+      "--iree-input-type=stablehlo_xla",
+      "--iree-execution-model=async-internal"
+    ]
+
+    {:ok, %{bytecode: %NxIREE.Module{bytecode: bytecode}, output_container: output_container}} =
+      NxIREE.Compiler.to_bytecode(fun, input_templates, iree_compiler_flags: compiler_flags)
+
+    socket =
+      socket
+      |> assign(:bytecode, Base.encode64(bytecode))
+      |> assign(:output_container, output_container)
+      |> assign(:function_signature, get_signature(function, input_templates, output_container))
+
+    {:noreply, socket}
+  end
+
+  defp get_signature({mod, fun, _a}, input_templates, output_container) do
+    "#{inspect(mod)}.#{fun}(#{to_flat_type(input_templates)}) -> #{to_flat_type(output_container)}"
+  end
+
+  defp to_flat_type(container) do
+    List.wrap(container)
+    |> Nx.Defn.Composite.flatten_list()
+    |> Enum.map(fn t ->
+      type_as_string(t) <> "x" <> Enum.join(Tuple.to_list(Nx.shape(t)), "x")
+    end)
+    |> Enum.join(", ")
+  end
+
+  defp type_as_string(tensor) do
+    {t, s} = Nx.type(tensor)
+
+    "#{t}#{s}"
+  end
 end

liveview_native/live_nx_iree/lib/live_nx_iree_web/live/swiftui/home_live.swiftui.neex

@@ -1 +1 @@
-<Text>Hello, World</Text>
+<NxFunction signature={@function_signature} bytecode={@bytecode}/>

liveview_native/live_nx_iree/native/swiftui/LiveNxIREE.xcodeproj/project.xcworkspace/xcuserdata/paulo.valente.xcuserdatad/UserInterfaceState.xcuserstate

@@ -8,10 +8,16 @@
 import SwiftUI
 import LiveViewNative
 
+@LiveElement
 struct NxFunctionView<Root: RootRegistry>: View {
+    @_documentation(visibility: public)
+    @LiveAttribute("bytecode") private var bytecode: String? = nil
+    @LiveAttribute("signature") private var signature: String? = nil
+    
     var body: some View {
-        Text("NxFunction Component")
-            .padding()
+        if signature != nil {
+            Text(signature!)
+                .padding()
+        }
     }
 }
-