updated sharding implementation to be inside the existing compiler and defn files. only the Mesh module has a standalone file now

Author: Chapaman

nx/lib/nx/defn.ex

@@ -849,6 +849,29 @@ defmodule Nx.Defn do
     :ok
   end
 
+  def shard_jit(fun, mesh, opts \\ []) when is_function(fun) and is_list(opts) do
+    wrap(fun, &shard_jit_apply(fun, mesh, &1, opts))
+  end
+
+  def shard_jit_apply(fun, mesh, args, opts \\ [])
+    when is_function(fun) and is_list(args) and is_list(opts) do
+  {on_conflict, opts} = Keyword.pop(opts, :on_conflict, :raise)
+
+    cond do
+      Nx.Defn.current() == nil ->
+        do_shard_jit_apply(fun, mesh, args, opts)
+
+      on_conflict == :raise ->
+        raise "cannot invoke Shard JITed function when there is a Shard JIT compilation happening"
+
+      on_conflict == :force ->
+        do_shard_jit_apply(fun, mesh, args, opts)
+
+      on_conflict == :reuse ->
+        apply(fun, args)
+    end
+end
+
   defp compile_error!(env, description) do
     raise CompileError, line: env.line, file: env.file, description: description
   end

nx/lib/nx/defn/compiler.ex

@@ -73,6 +73,25 @@ defmodule Nx.Defn.Compiler do
   """
   @callback __to_backend__(keyword) :: {module, keyword}
 
+  @doc """
+  Callback for compilation.
+
+  Its main purpose is to compile a function for a given mesh.
+  """
+  @callback __shard_jit__(
+    key :: term,
+    mesh :: Nx.Defn.Shard.Mesh.t(),
+    [vars],
+    fun :: (vars -> Nx.Container.t()),
+    args_list :: [[(-> Nx.Tensor.t())]],
+    opts :: keyword
+  ) :: [Nx.Container.t()]
+  when vars: [Nx.Container.t()]
+
+  @optional_callbacks [
+    __shard_jit__: 6
+  ]
+
   # Modules allowed in defn
   @allowed_modules [Nx.Constants, Nx.Defn, Nx.Defn.Kernel, Nx.LinAlg, Nx.Type]
 
@@ -265,6 +284,14 @@ defmodule Nx.Defn.Compiler do
     {:__block__, [], quoted}
   end
 
+  def __shard_jit__(fun, mesh, params, args_list, opts) do
+    {module, runtime_fun, opts} = prepare_options(fun, mesh, opts)
+    module.__shard_jit__(fun, mesh, params, runtime_fun, args_list, opts)
+  rescue
+    e in [UndefinedFunctionError] ->
+      raise_missing_callback(e, :__shard_jit__, 6, __STACKTRACE__)
+  end
+
   defp compile_prepare_arities(definitions) do
     for {{name, arity}, %{defaults: defaults}} <- definitions,
         arity <- (arity - map_size(defaults))..arity,

nx/lib/nx/defn/shard.ex

@@ -0,0 +1,4 @@
+defmodule Nx.Defn.Shard.Mesh do
+  defstruct [:name, :shape]
+  @type t :: %__MODULE__{name: String.t(), shape: tuple()}
+end