feat: add EXLA.to_mlir_module/2

exla/lib/exla.ex

@@ -355,6 +355,43 @@ defmodule EXLA do
     Nx.Defn.stream(function, args, Keyword.put(options, :compiler, EXLA))
   end
 
+  @doc """
+  Takes in a function, the templates variables and the compilation options
+  and returns the `EXLA.Executable` struct.
+
+  ## Examples
+
+      iex> fun = fn x, y -> Nx.add(Nx.sin(x), Nx.cos(y)) end
+      iex> args = [1.0, 2.0]
+      iex> module = EXLA.to_mlir_module(fun, args)
+      iex> EXLA.MLIR.Module.as_string(module)
+      ~c\"\"\"
+      module {
+        func.func public @main(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
+          %0 = stablehlo.sine %arg0 : tensor<f32>
+          %1 = stablehlo.cosine %arg1 : tensor<f32>
+          %2 = stablehlo.add %0, %1 : tensor<f32>
+          stablehlo.return %2 : tensor<f32>
+        }
+      }
+      \"\"\"
+  """
+  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})
+    end
+
+    function
+    |> jit([
+      {EXLA, {&EXLA.Defn.LockedCache.run/2, comp_fun}},
+      {:module_compilation, :to_mlir} | options
+    ])
+    |> apply(args)
+  catch
+    {:mlir_module, ref} -> %EXLA.MLIR.Module{ref: ref}
+  end
+
   @doc """
   Checks if the compilation of function with args is cached.
 

exla/lib/exla/mlir/module.ex

@@ -71,6 +71,11 @@ defmodule EXLA.MLIR.Module do
     * `:use_spmd` - enables Single-Program Multiple-Data partioning.
       This is set to true if `:num_partitions` is more than one, otherwise is `false`.
 
+    * `:module_compilation` - either `:to_mlir` or `:to_pjrt`. The default is `:to_pjrt`.
+
+      * `:to_pjrt` - the `EXLA.Executable` `:ref` field will hold the reference to a PjRt executable.
+      * `:to_mlir` - the `EXLA.Executable` `:ref` field will hold the reference to an MLIR module.
+
   Currently those options do not have an effect as they related to running the
   same compiled executable on multiple replicas.
 
@@ -102,16 +107,22 @@ defmodule EXLA.MLIR.Module do
     # module |> as_string() |> IO.puts()
 
     ref =
-      EXLA.NIF.mlir_compile(
-        client.ref,
-        module.ref,
-        Enum.map(argument_typespecs, &EXLA.Typespec.nif_encode/1),
-        num_replicas,
-        num_partitions,
-        use_spmd,
-        device_id
-      )
-      |> unwrap!()
+      case Keyword.get(options, :module_compilation, :to_pjrt) do
+        :to_mlir ->
+          module.ref
+
+        :to_pjrt ->
+          EXLA.NIF.mlir_compile(
+            client.ref,
+            module.ref,
+            Enum.map(argument_typespecs, &EXLA.Typespec.nif_encode/1),
+            num_replicas,
+            num_partitions,
+            use_spmd,
+            device_id
+          )
+          |> unwrap!()
+      end
 
     %Executable{
       client: client,