DRAFT: Add sharding support for data-parallelism

build.sbt

@@ -84,6 +84,8 @@ lazy val core = (project in file("core"))
     fork := true,
     javaOptions ++= scalapyJavaOptions,
     Test / envVars += "DIMWIT_SKIP_SYNC" -> "true",
+    // The sharding tests need a mesh of devices; a default JAX process reports a single CPU device.
+    Test / envVars += "XLA_FLAGS" -> "--xla_force_host_platform_device_count=8",
     coverageMinimumStmtTotal := 80,
     coverageFailOnMinimum := false,
     coverageHighlighting := true,

core/src/main/resources/python/jax_helper.py

@@ -85,3 +85,27 @@ def jit(f):
 
 def jit_fn(f, jit_kwargs=None):
     return wrap(jax.jit, f, kwargs=jit_kwargs)
+
+# --- sharding -------------------------------------------------------------
+
+def device_mesh(devices, shape, axis_names):
+    """Builds a jax.sharding.Mesh from a flat, row-major list of devices."""
+    import numpy as np
+    grid = np.array(list(devices), dtype=object).reshape(tuple(shape))
+    return jax.sharding.Mesh(grid, axis_names=tuple(axis_names))
+
+def named_sharding(mesh, rank, axis_index, mesh_axis_name):
+    """NamedSharding splitting axis `axis_index` of a rank-`rank` array over `mesh_axis_name`.
+
+    Every other axis is replicated.
+    """
+    spec = [None] * rank
+    spec[axis_index] = mesh_axis_name
+    return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(*spec))
+
+def replicate(x):
+    """Gathers a sharded array so that every device of its mesh holds the whole thing."""
+    sharding = x.sharding
+    if not hasattr(sharding, "mesh"):
+        return x
+    return jax.device_put(x, jax.sharding.NamedSharding(sharding.mesh, jax.sharding.PartitionSpec()))

core/src/main/scala/dimwit/package.scala

@@ -39,6 +39,16 @@ package object dimwit:
     given [A, B](using labelA: Label[A], labelB: Label[B]): Label[A |+| B] with
       val name: String = s"${labelA.name}+${labelB.name}"
 
+  /** Marks the data axis `A` as sharded over the mesh axis `M`, as in `Tensor2[Batch |@| X, Feature, Float32]`.
+    *
+    * `Batch |@| X` is just another axis label, so every tensor operation applies unchanged.
+    */
+  @targetName("Sharded")
+  infix trait |@|[A, M]
+  object `|@|`:
+    given [A, M](using label: Label[A], meshLabel: MeshLabel[M]): Label[A |@| M] with
+      val name: String = s"${label.name}@${meshLabel.name}"
+
   // Export tensor and related types
   export dimwit.tensor.{Tensor, Tensor0, Tensor1, Tensor2, Tensor3, Tensor4, TypedIndex}
   export dimwit.tensor.{Shape, Shape0, Shape1, Shape2, Shape3}
@@ -88,6 +98,10 @@ package object dimwit:
   export dimwit.stats.{Prob, LogProb}
   export dimwit.stats.{Distribution, IndependentDistribution, MultivariateDistribution, UnivariateDistribution}
 
+  // Export sharding stuff
+  export dimwit.sharding.{Mesh, Mesh1, Mesh2, Mesh3, MeshAxis, MeshAxisExtent, MeshAxisIndex, MeshLabel, MeshLabels}
+  export dimwit.sharding.ShardingOps.*
+
   /** Memory management helpfer making sure
     * all python objects allocated ar freed
     * after the function is executed.

core/src/main/scala/dimwit/sharding/Mesh.scala

@@ -0,0 +1,83 @@
+package dimwit.sharding
+
+import dimwit.hardware.Device
+import dimwit.jax.Jax
+import me.shadaj.scalapy.py.SeqConverters
+
+import scala.annotation.implicitNotFound
+
+/** A grid of devices, binding each mesh axis label of `M` to a number of devices. */
+final class Mesh[M <: Tuple: MeshLabels] private[sharding] (
+    val axisSizes: List[Int],
+    val devices: Seq[Device]
+):
+
+  lazy val axisNames: List[String] = MeshLabels[M].names
+
+  def size: Int = axisSizes.product
+
+  def sizeOf[A](meshAxis: MeshAxis[A])(using ev: MeshAxisIndex[M, A]): Int = axisSizes(ev.index)
+
+  private[dimwit] lazy val jaxMesh: Jax.PyDynamic =
+    Jax.jax_helper
+      .device_mesh(devices.map(_.toJaxDevice).toPythonProxy, axisSizes.toPythonProxy, axisNames.toPythonProxy)
+      .as[Jax.PyDynamic]
+
+  override def toString: String =
+    axisNames.zip(axisSizes).map((name, size) => s"$name -> $size").mkString("Mesh(", ", ", ")")
+
+object Mesh:
+
+  private[sharding] type ExtractLabel[Extent] = Extent match
+    case MeshAxisExtent[a] => a
+
+  private[sharding] type ExtractLabels[Extents <: Tuple] = Tuple.Map[Extents, ExtractLabel]
+
+  def apply[A: MeshLabel](extent: MeshAxisExtent[A]): Mesh1[A] = Mesh1(extent)
+  def apply[A: MeshLabel, B: MeshLabel](extent1: MeshAxisExtent[A], extent2: MeshAxisExtent[B]): Mesh2[A, B] = Mesh2(extent1, extent2)
+  def apply[A: MeshLabel, B: MeshLabel, C: MeshLabel](extent1: MeshAxisExtent[A], extent2: MeshAxisExtent[B], extent3: MeshAxisExtent[C]): Mesh3[A, B, C] = Mesh3(extent1, extent2, extent3)
+  def apply[Extents <: Tuple](extents: Extents)(using MeshLabels[ExtractLabels[Extents]]): Mesh[ExtractLabels[Extents]] = fromTuple(extents)
+
+  def fromTuple[Extents <: Tuple](extents: Extents)(using labels: MeshLabels[ExtractLabels[Extents]]): Mesh[ExtractLabels[Extents]] =
+    val sizes = extents.toList.collect:
+      case extent: MeshAxisExtent[?] => extent.size
+    val required = sizes.product
+    val available = Jax.devices
+    require(
+      available.size >= required,
+      s"Mesh ${labels.names.zip(sizes).map((name, size) => s"$name -> $size").mkString("(", ", ", ")")} requires $required devices, but only ${available.size} are available"
+    )
+    new Mesh(sizes, available.take(required))
+
+type Mesh1[A] = Mesh[Tuple1[A]]
+type Mesh2[A, B] = Mesh[(A, B)]
+type Mesh3[A, B, C] = Mesh[(A, B, C)]
+
+object Mesh1:
+  def apply[A: MeshLabel](extent: MeshAxisExtent[A]): Mesh1[A] = Mesh.fromTuple(Tuple1(extent))
+
+object Mesh2:
+  def apply[A: MeshLabel, B: MeshLabel](extent1: MeshAxisExtent[A], extent2: MeshAxisExtent[B]): Mesh2[A, B] =
+    Mesh.fromTuple((extent1, extent2))
+
+object Mesh3:
+  def apply[A: MeshLabel, B: MeshLabel, C: MeshLabel](
+      extent1: MeshAxisExtent[A],
+      extent2: MeshAxisExtent[B],
+      extent3: MeshAxisExtent[C]
+  ): Mesh3[A, B, C] = Mesh.fromTuple((extent1, extent2, extent3))
+
+/** Finds the position of a mesh axis in a mesh. */
+@implicitNotFound("MeshAxis[${A}] not found in Mesh[${M}]")
+trait MeshAxisIndex[M <: Tuple, A]:
+  def index: Int
+
+object MeshAxisIndex:
+
+  def apply[M <: Tuple, A](using idx: MeshAxisIndex[M, A]): Int = idx.index
+
+  given found[A, Tail <: Tuple]: MeshAxisIndex[A *: Tail, A] with
+    val index = 0
+
+  given search[H, T <: Tuple, A](using next: MeshAxisIndex[T, A]): MeshAxisIndex[H *: T, A] with
+    val index = 1 + next.index

core/src/main/scala/dimwit/sharding/MeshAxis.scala

@@ -0,0 +1,14 @@
+package dimwit.sharding
+
+/** Represents an axis of a device [[Mesh]]. The mesh-side counterpart of [[dimwit.tensor.Axis]]. */
+final class MeshAxis[A: MeshLabel]:
+
+  def name: String = summon[MeshLabel[A]].name
+
+  def extent(size: Int): MeshAxisExtent[A] = MeshAxisExtent(this, size)
+  def ->(size: Int): MeshAxisExtent[A] = this.extent(size)
+
+  override def toString: String = s"MeshAxis($name)"
+
+/** A mesh axis together with the number of devices along it. */
+case class MeshAxisExtent[A: MeshLabel](axis: MeshAxis[A], size: Int)

core/src/main/scala/dimwit/sharding/MeshLabels.scala

@@ -0,0 +1,51 @@
+package dimwit.sharding
+
+import scala.quoted.*
+
+/** A label for an axis of a device [[Mesh]], declared with `trait X derives MeshLabel`.
+  * The mesh-side counterpart of [[dimwit.tensor.Label]].
+  */
+@scala.annotation.implicitNotFound("""
+A mesh axis label ${T} was given or inferred, which does not have a MeshLabel instance.
+Ensure that all mesh axis types ${T} are defined with 'derives MeshLabel' (e.g. 'trait X derives MeshLabel')
+""")
+trait MeshLabel[T]:
+  def name: String
+
+object MeshLabel:
+
+  def apply[T](using meshLabel: MeshLabel[T]): MeshLabel[T] = meshLabel
+
+  inline def derived[T]: MeshLabel[T] = ${ derivedMacro[T] }
+
+  private def derivedMacro[T: Type](using Quotes): Expr[MeshLabel[T]] =
+    import quotes.reflect.*
+    val tpe = TypeRepr.of[T]
+    val simpleName = tpe.typeSymbol.name
+    '{
+      new MeshLabel[T]:
+        def name: String = ${ Expr(simpleName) }
+    }
+
+@scala.annotation.implicitNotFound("""
+A tuple of mesh axis labels ${T} was given or inferred that does not have a valid MeshLabels instance.
+
+Ensure that all of the types in the tuple have a 'derives MeshLabel' clause.
+""")
+trait MeshLabels[T]:
+  def names: List[String]
+
+private class MeshLabelsImpl[T](val names: List[String]) extends MeshLabels[T]
+
+object MeshLabels:
+
+  def apply[T](using labels: MeshLabels[T]): MeshLabels[T] = labels
+
+  given emptyTuple: MeshLabels[EmptyTuple] = new MeshLabelsImpl[EmptyTuple](Nil)
+
+  given lift[A](using v: MeshLabel[A]): MeshLabels[A] = new MeshLabelsImpl[A](List(v.name))
+
+  given consTuple[H, T <: Tuple](using
+      head: MeshLabel[H],
+      tail: MeshLabels[T]
+  ): MeshLabels[H *: T] = new MeshLabelsImpl[H *: T](head.name :: tail.names)

core/src/main/scala/dimwit/sharding/ShardingOps.scala

@@ -0,0 +1,51 @@
+package dimwit.sharding
+
+import dimwit.jax.Jax
+import dimwit.|@|
+import dimwit.tensor.Axis
+import dimwit.tensor.Label
+import dimwit.tensor.Labels
+import dimwit.tensor.ShapeTypeHelpers.AxisReplacer
+import dimwit.tensor.Tensor
+
+object ShardingOps:
+
+  extension [T <: Tuple: Labels, V](t: Tensor[T, V])
+
+    /** Splits `t` across `mesh` along one axis, placing one shard per device and rewriting that axis from `L` to `L |@| A`.
+      *
+      * {{{
+      * val mesh = Mesh1(MeshAxis[X] -> 4)
+      * val sharded: Tensor2[Batch |@| X, Feature, Float32] = t.shard(mesh, Axis[Batch] -> MeshAxis[X])
+      * }}}
+      *
+      * @throws IllegalArgumentException if the extent of `L` is not divisible by the mesh axis size.
+      */
+    def shard[M <: Tuple, L: Label, A: MeshLabel](mesh: Mesh[M], mapping: (Axis[L], MeshAxis[A]))(using
+        replacer: AxisReplacer[T, L, L |@| A],
+        meshIndex: MeshAxisIndex[M, A],
+        labels: Labels[replacer.NewShape]
+    ): Tensor[replacer.NewShape, V] =
+      val axisName = summon[Label[L]].name
+      val meshAxisName = mapping._2.name
+      val extent = t.shape.dimensions(replacer.index)
+      val meshSize = mesh.axisSizes(meshIndex.index)
+      if extent % meshSize != 0 then
+        throw new IllegalArgumentException(
+          s"Cannot shard axis $axisName of extent $extent over mesh axis $meshAxisName of size $meshSize: " +
+            s"$extent is not divisible by $meshSize."
+        )
+      val sharding = Jax.jax_helper.named_sharding(mesh.jaxMesh, t.shape.rank, replacer.index, meshAxisName)
+      Tensor[replacer.NewShape, V](Jax.device_put(t.jaxValue, sharding.as[Jax.PyDynamic]))
+
+    /** Gathers a sharded axis back across the mesh, rewriting it from `L |@| A` to `L` so that every device holds the whole tensor.
+      *
+      * {{{
+      * val gathered: Tensor2[Batch, Feature, Float32] = sharded.unshard(Axis[Batch |@| X])
+      * }}}
+      */
+    def unshard[L, A](axis: Axis[L |@| A])(using
+        replacer: AxisReplacer[T, L |@| A, L],
+        labels: Labels[replacer.NewShape]
+    ): Tensor[replacer.NewShape, V] =
+      Tensor[replacer.NewShape, V](Jax.jax_helper.replicate(t.jaxValue).as[Jax.PyDynamic])