fix: add support for defn compilation in EXLA.to_mlir_module

exla/lib/exla.ex

@@ -360,11 +360,18 @@ defmodule EXLA do
   Takes in a function, the argument templates and the compilation
   options and returns the textual representation of the MLIR module.
 
+  ## Options
+
+    * `:nested_defn_compilation` - a boolean that indicates whether
+      this function is being called from within a `defn` compiler.
+      Defaults to `false`.
+
   ## Examples
 
       iex> fun = fn x, y -> Nx.add(Nx.sin(x), Nx.cos(y)) end
       iex> args = [1.0, 2.0]
-      iex> EXLA.to_mlir_module(fun, args)
+      iex> %{mlir_module: mlir_module} = EXLA.to_mlir_module(fun, args)
+      iex> mlir_module
       """
       module {
         func.func public @main(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
@@ -378,19 +385,31 @@ defmodule EXLA do
   '''
   def to_mlir_module(function, args, options \\ []) do
     comp_fun = fn _key, callback ->
-      {:ok, {_xla_time, executable, _extra, _outfeed}} = callback.()
-      throw({:mlir_module, executable.ref})
+      {:ok, {_xla_time, executable, {_, used_inputs, outputs}, _outfeed}} = callback.()
+      throw({:mlir_module, executable.ref, used_inputs, outputs})
     end
 
-    opts = [
-      {EXLA, {&EXLA.Defn.LockedCache.run/2, comp_fun}},
-      {:module_compilation, :to_mlir} | options
-    ]
+    {nested_compilation?, options} = Keyword.pop(options, :nested_defn_compilation, false)
 
-    jit_apply(function, args, opts)
+    opts =
+      Keyword.merge(options, [
+        {EXLA, {&EXLA.Defn.LockedCache.run/2, comp_fun}},
+        module_compilation: :to_mlir,
+        compiler: EXLA
+      ])
+
+    if nested_compilation? do
+      EXLA.Defn.__compile__(function, args, function, opts)
+    else
+      Nx.Defn.compile(function, args, opts)
+    end
   catch
-    {:mlir_module, ref} ->
-      EXLA.MLIR.Module.as_string(%EXLA.MLIR.Module{ref: ref})
+    {:mlir_module, ref, used_inputs, output_container} ->
+      %{
+        used_inputs: used_inputs,
+        output_container: output_container,
+        mlir_module: EXLA.MLIR.Module.as_string(%EXLA.MLIR.Module{ref: ref})
+      }
   end
 
   @doc """

exla/lib/exla/defn.ex

@@ -406,7 +406,8 @@ defmodule EXLA.Defn do
       outfeed = Outfeed.new(hooks, defined_hooks)
       comp_key = {ref, client.name, outfeed.used_hooks, lazy_transfers, options}
 
-      {comp_time, {evaled, {xla_time, executable, inputs_and_typespecs, outfeed}}} =
+      {comp_time,
+       {evaled, {xla_time, executable, {inputs_and_typespecs, _used_inputs, _outputs}, outfeed}}} =
         :timer.tc(fn ->
           comp_cache_fun.(comp_key, fn ->
             {reverse_inputs_and_typespecs, reverse_infeeds} =
@@ -466,7 +467,9 @@ defmodule EXLA.Defn do
                   )
                 end)
 
-              {:ok, {xla_time, executable, inputs_and_typespecs, %{outfeed | infeeds: []}}}
+              {:ok,
+               {xla_time, executable, {inputs_and_typespecs, used_inputs, outputs},
+                %{outfeed | infeeds: []}}}
             end)
           end)
         end)

exla/test/exla_test.exs

@@ -17,4 +17,56 @@ defmodule EXLATest do
       end
     end
   end
+
+  defmodule ValidCompiler do
+    def __jit__(key, vars, fun, args_list, opts) do
+      __compile__(key, vars, fun, opts).(args_list)
+    end
+
+    def __compile__(_key, vars, fun, opts) do
+      result = EXLA.to_mlir_module(fun, vars, Keyword.put(opts, :nested_defn_compilation, true))
+      throw({__MODULE__, result})
+    end
+  end
+
+  defmodule InvalidCompiler do
+    def __jit__(key, vars, fun, args_list, opts) do
+      __compile__(key, vars, fun, opts).(args_list)
+    end
+
+    def __compile__(_key, vars, fun, opts) do
+      EXLA.to_mlir_module(fun, vars, Keyword.put(opts, :nested_defn_compilation, false))
+    end
+  end
+
+  describe "to_mlir_module/3" do
+    test "fails if the compiler doesn't set the nested compilation flag" do
+      assert_raise BadArityError, fn ->
+        Nx.Defn.jit_apply(&Nx.add/2, [1, 2], compiler: __MODULE__.InvalidCompiler)
+      end
+    end
+
+    test "works if the compiler sets the nested compilation flag" do
+      try do
+        Nx.Defn.jit_apply(&Nx.add/2, [1, 2], compiler: __MODULE__.ValidCompiler)
+      catch
+        {__MODULE__.ValidCompiler, result} ->
+          assert %{mlir_module: module, output_container: container, used_inputs: used_inputs} =
+                   result
+
+          assert module == """
+                 module {
+                   func.func public @main(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
+                     %0 = stablehlo.add %arg0, %arg1 : tensor<i32>
+                     return %0 : tensor<i32>
+                   }
+                 }
+                 """
+
+          assert Nx.compatible?(container, Nx.template({}, :s32))
+
+          assert used_inputs == %{0 => nil, 1 => nil}
+      end
+    end
+  end
 end