Basic Nx.block implementation (#1709)

Co-authored-by: Paulo Valente <16843419+polvalente@users.noreply.github.com>

Author: Chapaman

exla/lib/exla/backend.ex

@@ -318,12 +318,12 @@ defmodule EXLA.Backend do
   end
 
   @impl true
-  def optional(name, args, fun) do
+  def block(struct, _output, args, fun) do
     # Here we take the leading tensor arguments and pass them as JIT arguments
     {tensors, rest} = Enum.split_while(args, &is_struct(&1, Nx.Tensor))
 
     wrapper_fun = fn tensors ->
-      Nx.Defn.Expr.optional(name, Tuple.to_list(tensors) ++ rest, fun)
+      Nx.Defn.Expr.block(struct, nil, Tuple.to_list(tensors) ++ rest, fun)
     end
 
     jit([], wrapper_fun, tensors, [List.to_tuple(tensors)])

exla/lib/exla/defn.ex

@@ -601,11 +601,12 @@ defmodule EXLA.Defn do
   end
 
   defp cached_recur_operator(
-         :optional,
+         :block,
          %T{
            data: %Expr{
              args: [
-               %{data: %{op: :qr, args: [tensor, _opts]}},
+               %Nx.Block.QR{},
+               [tensor],
                {%{type: {type_kind, _}} = q_expr, r_expr},
                _callback
              ]
@@ -631,11 +632,12 @@ defmodule EXLA.Defn do
   end
 
   defp cached_recur_operator(
-         :optional,
+         :block,
          %T{
            data: %Expr{
              args: [
-               %{data: %{op: :eigh, args: [tensor, _opts]}},
+               %Nx.Block.Eigh{},
+               [tensor],
                {%{type: {evec_type_kind, _}} = eigenvals_expr,
                 %{type: {eval_type_kind, _}} = eigenvecs_expr},
                _callback
@@ -672,16 +674,15 @@ defmodule EXLA.Defn do
   end
 
   defp cached_recur_operator(
-         :optional,
+         :block,
          %T{
            data: %Expr{
-             args: [%{data: %{op: :take, args: [tensor, indices, opts]}}, expr, _callback]
+             args: [%Nx.Block.Take{axis: axis}, [tensor, indices], expr, _callback]
            }
          },
          state,
          cache
        ) do
-    axis = opts[:axis]
     {tensor, cache} = recur_operator(tensor, state, cache) |> unwrap_single_tensor!()
     {indices, cache} = recur_operator(indices, state, cache) |> unwrap_single_tensor!()
 
@@ -714,56 +715,80 @@ defmodule EXLA.Defn do
   end
 
   defp cached_recur_operator(
-         :optional,
-         %T{data: %Expr{args: [%{data: %{op: :top_k, args: [tensor, opts]}}, expr, _callback]}},
+         :block,
+         %T{data: %Expr{args: [%Nx.Block.TopK{k: k}, [tensor], expr, _callback]}},
          state,
          cache
        ) do
     {tensor, cache} = recur_operator(tensor, state, cache) |> unwrap_single_tensor!()
     {values, idx} = expr
     typespecs = [expr_to_typespec(values), expr_to_typespec(idx)]
-    results = Value.top_k(tensor, opts[:k], typespecs)
+    results = Value.top_k(tensor, k, typespecs)
     {results, cache}
   end
 
   defp cached_recur_operator(
-         :optional,
-         %T{data: %Expr{args: [%{data: %{op: :fft2, args: [tensor, opts]}}, expr, _callback]}},
+         :block,
+         %T{data: %Expr{args: [%Nx.Block.FFT2{} = fft2_struct, [tensor], expr, _callback]}},
          state,
          cache
        ) do
     {tensor, cache} = recur_operator(tensor, state, cache) |> unwrap_single_tensor!()
 
+    opts = [lengths: fft2_struct.lengths, axes: fft2_struct.axes]
+
+    opts =
+      if eps = fft2_struct.eps do
+        Keyword.put(opts, :eps, eps)
+      else
+        opts
+      end
+
     {fft2(&Value.fft(&1, :fft, &2, &3), [tensor, opts], expr, state), cache}
   end
 
   defp cached_recur_operator(
-         :optional,
-         %T{data: %Expr{args: [%{data: %{op: :ifft2, args: [tensor, opts]}}, expr, _callback]}},
+         :block,
+         %T{data: %Expr{args: [%Nx.Block.IFFT2{} = ifft2_struct, [tensor], expr, _callback]}},
          state,
          cache
        ) do
     {tensor, cache} = recur_operator(tensor, state, cache) |> unwrap_single_tensor!()
 
+    opts = [lengths: ifft2_struct.lengths, axes: ifft2_struct.axes]
+
+    opts =
+      if eps = ifft2_struct.eps do
+        Keyword.put(opts, :eps, eps)
+      else
+        opts
+      end
+
     {fft2(&Value.fft(&1, :ifft, &2, &3), [tensor, opts], expr, state), cache}
   end
 
-  defp cached_recur_operator(:optional, %T{data: %Expr{args: args}}, state, cache) do
-    [call, expr, _callback] = args
-    %{data: %{args: in_args, op: op}} = call
-
-    {args, opts} = Enum.split_while(in_args, &(not is_list(&1)))
+  defp cached_recur_operator(:block, %T{data: %Expr{args: args}}, state, cache) do
+    [struct, in_args, expr, _callback] = args
+    %module{} = struct
 
-    {call_args, cache} = Enum.map_reduce(args, cache, &recur_operator(&1, state, &2))
-    key = computation_key(op, call_args ++ opts)
+    {call_args, cache} = Enum.map_reduce(in_args, cache, &recur_operator(&1, state, &2))
+    key = computation_key(module, [struct | call_args])
 
     {call_body, cache} =
       case cache do
         %{^key => computation} ->
           {computation, cache}
 
         %{} ->
-          {computation, cache} = optional_computation("optional", call_args, expr, state, cache)
+          {computation, cache} =
+            block_computation(
+              block_subfunction_description(struct),
+              call_args,
+              expr,
+              state,
+              cache
+            )
+
           {computation, Map.put(cache, key, computation)}
       end
 
@@ -1818,8 +1843,14 @@ defmodule EXLA.Defn do
     {region, merge_outfeed(cache, comp_cache)}
   end
 
-  defp optional_computation(name, args, expr, %{builder: %Function{}} = state, cache) do
-    %Function{module: module, name: name} = subbuilder(state.builder, name)
+  defp block_subfunction_description(%module{} = _) do
+    module
+    |> Atom.to_string()
+    |> String.replace(".", "_")
+  end
+
+  defp block_computation(description, args, expr, %{builder: %Function{}} = state, cache) do
+    %Function{module: module, name: name} = subbuilder(state.builder, description)
 
     arg_typespecs = Enum.map(args, &Value.get_typespec/1)
     out_typespecs = container_to_typespecs(expr)

exla/mix.lock

@@ -11,7 +11,6 @@
   "makeup_erlang": {:hex, :makeup_erlang, "1.0.3", "4252d5d4098da7415c390e847c814bad3764c94a814a0b4245176215615e1035", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "953297c02582a33411ac6208f2c6e55f0e870df7f80da724ed613f10e6706afd"},
   "nimble_parsec": {:hex, :nimble_parsec, "1.4.2", "8efba0122db06df95bfaa78f791344a89352ba04baedd3849593bfce4d0dc1c6", [:mix], [], "hexpm", "4b21398942dda052b403bbe1da991ccd03a053668d147d53fb8c4e0efe09c973"},
   "nimble_pool": {:hex, :nimble_pool, "1.0.0", "5eb82705d138f4dd4423f69ceb19ac667b3b492ae570c9f5c900bb3d2f50a847", [:mix], [], "hexpm", "80be3b882d2d351882256087078e1b1952a28bf98d0a287be87e4a24a710b67a"},
-  "nx": {:hex, :nx, "0.11.0", "d37723dbd6cfa274a5def6d6664f5680c32e2eb8a1ce25ec6d91751967fa0abf", [:mix], [{:complex, "~> 0.6", [hex: :complex, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "36157b21239aeb251d6cbac23eb0eb3495a5e1e0cbc2e6df16afd2ede1575205"},
   "statistex": {:hex, :statistex, "1.0.0", "f3dc93f3c0c6c92e5f291704cf62b99b553253d7969e9a5fa713e5481cd858a5", [:mix], [], "hexpm", "ff9d8bee7035028ab4742ff52fc80a2aa35cece833cf5319009b52f1b5a86c27"},
   "telemetry": {:hex, :telemetry, "1.4.1", "ab6de178e2b29b58e8256b92b382ea3f590a47152ca3651ea857a6cae05ac423", [:rebar3], [], "hexpm", "2172e05a27531d3d31dd9782841065c50dd5c3c7699d95266b2edd54c2dafa1c"},
   "xla": {:hex, :xla, "0.10.0", "41121e9f011456242d3a79b9289910ce43419be0b0e7ebe67cc1292c6b3f232f", [:make, :mix], [{:elixir_make, "~> 0.4", [hex: :elixir_make, repo: "hexpm", optional: false]}], "hexpm", "f57d91aea6e661b52bf12239316c598679e9170628122bbd941235f040122bc6"},

exla/test/exla/nx_linalg_doctest_test.exs

@@ -25,7 +25,7 @@ defmodule EXLA.NxLinAlgDoctestTest do
     least_squares: 3,
     determinant: 1,
     matrix_power: 2,
-    lu: 2,
+    lu: 1,
     qr: 2
   ]