Repository: google/jax
Pull request #23315 (2024-08-29 16:06:33)

Initial integration of sharding in types in JAX. Currently we just support `nary` ops in forward only sharding propagation. Currently this functionality is experimental and hidden behind `jax_sharding_in_types` config flag.
Initial integration of sharding in types in JAX. Currently we just support `nary` ops in forward only sharding propagation. Currently this functionality is experimental and hidden behind `jax_sharding_in_types` config flag.

There will be more improvements and semantics clarification coming in the future as we integrate it more into JAX.

Co-authored-by: Dougal Maclaurin <dougalm@google.com>



Patch:
diff --git a/jax/_src/array.py b/jax/_src/array.py
index 7659c180d..909f5acf0 100644
--- a/jax/_src/array.py
+++ b/jax/_src/array.py
@@ -41,7 +41,7 @@ from jax._src.lib import xla_client as xc
 from jax._src.lib import xla_extension as xe
 from jax._src.sharding import Sharding
 from jax._src.sharding_impls import (
-    PmapSharding, SingleDeviceSharding,
+    PmapSharding, SingleDeviceSharding, NamedSharding,
     device_replica_id_map, hashed_index, num_addressable_indices, local_to_global_shape)  # pyformat: disable
 from jax._src.typing import ArrayLike, DLDeviceType
 from jax._src.util import safe_zip, unzip3, use_cpp_class, use_cpp_method, cache
@@ -1012,7 +1012,13 @@ def make_array_from_single_device_arrays(
 core.pytype_aval_mappings[ArrayImpl] = abstract_arrays.canonical_concrete_aval
 xla.pytype_aval_mappings[ArrayImpl] = op.attrgetter('aval')
 xla.canonicalize_dtype_handlers[ArrayImpl] = pxla.identity
-api_util._shaped_abstractify_handlers[ArrayImpl] = op.attrgetter('aval')
+def _get_aval_array(self):
+  if config.sharding_in_types.value and isinstance(self.sharding, NamedSharding):
+    return self.aval.update(sharding=NamedSharding(
+        self.sharding.mesh.abstract_mesh, self.sharding.spec))
+  else:
+    return self.aval
+api_util._shaped_abstractify_handlers[ArrayImpl] = _get_aval_array
 # TODO(jakevdp) replace this with true inheritance at the C++ level.
 basearray.Array.register(ArrayImpl)
 
diff --git a/jax/_src/config.py b/jax/_src/config.py
index b6d2358f4..f17ca5938 100644
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -213,6 +213,7 @@ def trace_context():
           default_device.value, random_seed_offset.value,
           threefry_partitionable.value,
           threefry_gpu_kernel_lowering.value,
+          sharding_in_types.value,
           softmax_custom_jvp.value,
           enable_memories.value,
           disable_jit.value,
@@ -826,6 +827,7 @@ class _GlobalExtraJitContext(NamedTuple):
   random_seed_offset: int = 0
   threefry_partitionable: bool = False
   threefry_gpu_kernel_lowering: bool = False
+  sharding_in_types: bool = False
   softmax_custom_jvp: bool = False
   xla_profile_version: int = 0
   pgle_profiling_runs: int = 0
@@ -864,6 +866,7 @@ class _ThreadLocalExtraJitContext(NamedTuple):
   random_seed_offset: int | None = None
   threefry_partitionable: bool | None = None
   threefry_gpu_kernel_lowering: bool | None = None
+  sharding_in_types: bool | None = None
   softmax_custom_jvp: bool | None = None
   xla_profile_version: int | None = None
   pgle_profiling_runs: int | None = None
@@ -1139,6 +1142,16 @@ threefry_gpu_kernel_lowering = bool_state(
     update_thread_local_hook=lambda val: update_thread_local_jit_state(
         threefry_gpu_kernel_lowering=val))
 
+sharding_in_types = bool_state(
+    name='jax_sharding_in_types',
+    default=False,
+    help=('When True, enables forward only sharding propagation in JAX and '
+          'avals have sharding on them.'),
+    update_global_hook=lambda val: _update_global_jit_state(
+        sharding_in_types=val),
+    update_thread_local_hook=lambda val: update_thread_local_jit_state(
+        sharding_in_types=val))
+
 
 softmax_custom_jvp = bool_state(
     name='jax_softmax_custom_jvp',
diff --git a/jax/_src/core.py b/jax/_src/core.py
index f80cd0418..dbd4dc58e 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -1733,17 +1733,23 @@ def _invalid_shape_error(shape: Shape, context: str=""):
   return TypeError(msg)
 
 class ShapedArray(UnshapedArray):
-  __slots__ = ['shape']
+  __slots__ = ['shape', 'dtype', 'weak_type', 'sharding']
   array_abstraction_level = 2
   named_shape = {}  # type: ignore
 
-  def __init__(self, shape, dtype, weak_type=False, named_shape=None):
+  def __init__(self, shape, dtype, weak_type=False, named_shape=None,
+               sharding=None):
     del named_shape  # unused, vestigial
     self.shape = canonicalize_shape(shape)
     self.dtype = _dtype_object(dtype)
     self.weak_type = weak_type
+    if config.sharding_in_types.value:
+      self.sharding = sharding
+    else:
+      self.sharding = None
 
-  def update(self, shape=None, dtype=None, weak_type=None, named_shape=None):
+  def update(self, shape=None, dtype=None, weak_type=None, named_shape=None,
+             sharding=None):
     del named_shape  # unused, vestigial
     if shape is None:
       shape = self.shape
@@ -1751,7 +1757,9 @@ class ShapedArray(UnshapedArray):
       dtype = self.dtype
     if weak_type is None:
       weak_type = self.weak_type
-    return ShapedArray(shape, dtype, weak_type)
+    if sharding is None:
+      sharding = self.sharding
+    return ShapedArray(shape, dtype, weak_type, sharding=sharding)
 
   ndim = property(lambda self: len(self.shape))
   size = property(lambda self:
@@ -1766,13 +1774,14 @@ class ShapedArray(UnshapedArray):
   def __eq__(self, other):
     return (type(self) is type(other)
             and self.dtype == other.dtype and self.shape == other.shape
-            and self.weak_type == other.weak_type)
+            and self.weak_type == other.weak_type
+            and self.sharding == other.sharding)
 
   def __hash__(self):
     # can use hash(self.dtype) and rely on the fact that numpy reuses base dtype
     # objects, e.g. `np.zeros(3).dtype is np.zeros(4).dtype`, or we can use
     # the unique character code via hash(self.dtype.char)
-    return hash((self.shape, self.dtype, self.weak_type))
+    return hash((self.shape, self.dtype, self.weak_type, self.sharding))
 
   def at_least_vspace(self):
     return ShapedArray(self.shape, primal_dtype_to_tangent_dtype(self.dtype),
@@ -1791,7 +1800,10 @@ class ShapedArray(UnshapedArray):
     dt_str =  _short_dtype_name(self.dtype) if short_dtypes else self.dtype.name
     dt_str = dt_str.replace('void', 'float0')
     shapestr = ','.join(map(str, self.shape))
-    return f'{dt_str}[{shapestr}]'
+    if self.sharding is None:
+      return f'{dt_str}[{shapestr}]'
+    else:
+      return f'{dt_str}[{shapestr}]({self.sharding})'
 
   def _len(self, ignored_tracer):
     try:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index e81c6c157..29fa41012 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -65,7 +65,7 @@ from jax._src.lib import xla_client
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import chlo
 from jax._src.lib.mlir.dialects import hlo
-from jax._src.sharding_impls import PmapSharding
+from jax._src.sharding_impls import PmapSharding, NamedSharding, PartitionSpec
 from jax._src.typing import Array, ArrayLike, DimSize, DuckTypedArray, DTypeLike, Shape
 from jax._src.util import (cache, safe_zip, safe_map, canonicalize_axis,
                            split_list, NumpyComplexWarning)
@@ -1709,13 +1709,54 @@ def broadcasting_shape_rule(name, *avals):
   return tuple(result_shape)
 
 
+def broadcasting_sharding_rule(name, *avals):
+  shapes = [aval.shape for aval in avals if aval.shape]
+  if not shapes:
+    return ()
+  if len({len(shape) for shape in shapes}) != 1:
+    msg = '{}: arrays must have same number of dimensions, got {}.'
+    raise TypeError(msg.format(name, ', '.join(map(str, map(tuple, shapes)))))
+
+  specs = [a.sharding.spec for a in avals if a.shape]
+
+  mesh = None
+  for a in avals:
+    if a.shape:
+      mesh = a.sharding.mesh
+      if mesh is not None and mesh != a.sharding.mesh:
+        raise ValueError(
+            f'Mesh for all inputs should be equal. Got one mesh: {mesh} and'
+            f' another mesh: {a.sharding.mesh}')
+  assert mesh is not None
+
+  result_specs = []
+  for ss, ds in zip(zip(*specs), zip(*shapes)):
+    if all(s == ss[0] for s in ss[1:]):
+      # if all dimension shardings are same, the resulting dimension sharding is
+      # the same.
+      result_specs.append(ss[0])
+    else:
+      non_trivial_s = [s for s, d in zip(ss, ds)
+                       if not (core.definitely_equal(d, 1) and s is None)]
+      if not non_trivial_s:
+        result_specs.append(None)
+      elif all(non_trivial_s[0] == s for s in non_trivial_s[1:]):
+        result_specs.append(non_trivial_s[0])
+      else:
+        raise TypeError(f'{name} got incompatible shardings for broadcasting: '
+                        f'{", ".join(map(str, map(tuple, specs)))}.')
+  return NamedSharding(mesh, PartitionSpec(*result_specs))
+
+
 def naryop(result_dtype, accepted_dtypes, name, allow_extended_dtype=False,
            require_same_dtypes=False):
   dtype_rule = partial(naryop_dtype_rule, result_dtype, accepted_dtypes, name,
                        allow_extended_dtype=allow_extended_dtype,
                        require_same=require_same_dtypes)
   shape_rule = partial(broadcasting_shape_rule, name)
-  prim = standard_primitive(shape_rule, dtype_rule, name)
+  sharding_rule = partial(broadcasting_sharding_rule, name)
+  prim = standard_primitive(shape_rule, dtype_rule, name,
+                            sharding_rule=sharding_rule)
   batching.defbroadcasting(prim)
   pe.def_trivial_padding(prim)
   return prim
@@ -1772,6 +1813,20 @@ def broadcast_hlo(
     out.append(arg)
   return out
 
+def multi_sharding_in_dim(ctx, ops, in_avals, out_aval):
+  out = []
+  for op, in_aval in zip(ops, in_avals):
+    if in_aval.sharding == out_aval.sharding or in_aval.sharding is None:
+      out.append(op)
+    else:
+      # TODO(yashkatariya, dougalm): If `in_aval.sharding` contains
+      # CompilerShardingAxis, then specify `unspecified_dims` via
+      # `wrap_with_sharding_op`.
+      sp = in_aval.sharding._to_xla_hlo_sharding(in_aval.ndim).to_proto()
+      out.append(mlir.wrap_with_sharding_op(ctx, op, out_aval, sp))
+  return out
+
+
 def _nary_lower_hlo(op: Callable, ctx,
                     *args: ir.Value,
                     explicit_type=False, **params) -> Sequence[ir.Value]:
@@ -1782,13 +1837,19 @@ def _nary_lower_hlo(op: Callable, ctx,
   """
   del params
   avals_in, (aval_out,) = ctx.avals_in, ctx.avals_out
-  broadcasted_args = mlir.multi_broadcast_in_dim(
-      ctx, args, avals_in, aval_out.shape)
+  args = mlir.multi_broadcast_in_dim(ctx, args, avals_in, aval_out.shape)  # type: ignore
+  if config.sharding_in_types.value:
+    args = multi_sharding_in_dim(ctx, args, avals_in, aval_out)
 
   if explicit_type:
-    return [op(mlir.aval_to_ir_type(aval_out), *broadcasted_args)]
+    out = op(mlir.aval_to_ir_type(aval_out), *args)
   else:
-    return [op(*broadcasted_args)]
+    out = op(*args)
+  if config.sharding_in_types.value:
+    out_sp = aval_out.sharding._to_xla_hlo_sharding(aval_out.ndim).to_proto()
+    return [mlir.wrap_with_sharding_op(ctx, out, aval_out, out_sp)]
+  else:
+    return [out]
 
 
 _float = {np.floating}
@@ -2445,6 +2506,10 @@ def _convert_element_type_shape_rule(operand, *, new_dtype, weak_type,
                                      sharding):
   return operand.shape
 
+def _convert_element_type_sharding_rule(operand, *, new_dtype, weak_type,
+                                        sharding):
+  return sharding
+
 def _convert_element_type_dtype_rule(operand, *, new_dtype, weak_type,
                                      sharding):
   if (operand.dtype != new_dtype and
@@ -2538,7 +2603,8 @@ convert_element_type_p.def_impl(partial(dispatch.apply_primitive, convert_elemen
 convert_element_type_p.def_abstract_eval(
     partial(standard_abstract_eval, convert_element_type_p,
             _convert_element_type_shape_rule, _convert_element_type_dtype_rule,
-            _convert_element_type_weak_type_rule))
+            _convert_element_type_weak_type_rule,
+            _convert_element_type_sharding_rule))
 ad.defjvp(convert_element_type_p, _convert_element_type_jvp_rule)
 ad.primitive_transposes[convert_element_type_p] = _convert_element_type_transpose_rule
 batching.defvectorized(convert_element_type_p)
diff --git a/jax/_src/lax/utils.py b/jax/_src/lax/utils.py
index 01301db1a..5e3e9bcd8 100644
--- a/jax/_src/lax/utils.py
+++ b/jax/_src/lax/utils.py
@@ -20,6 +20,7 @@ from functools import partial
 
 from jax._src import core
 from jax._src import dispatch
+from jax._src import config
 from jax._src import dtypes
 from jax._src.util import safe_zip
 from jax._src.lib import xla_client
@@ -37,19 +38,19 @@ def _argnum_weak_type(*argnums):
   return lambda *args, **_: all(args[i].weak_type for i in argnums)
 
 def standard_primitive(shape_rule, dtype_rule, name,
-                       weak_type_rule=None):
+                       weak_type_rule=None, sharding_rule=None):
   weak_type_rule = weak_type_rule or _standard_weak_type_rule
   prim = core.Primitive(name)
   prim.def_impl(partial(dispatch.apply_primitive, prim))
   prim.def_abstract_eval(
       partial(standard_abstract_eval, prim, shape_rule, dtype_rule,
-              weak_type_rule))
+              weak_type_rule, sharding_rule))
   return prim
 
 def _get_array_abstraction_level(a): return a.array_abstraction_level
 
 def standard_abstract_eval(prim, shape_rule, dtype_rule, weak_type_rule,
-                           *avals, **kwargs):
+                           sharding_rule, *avals, **kwargs):
   assert all(isinstance(aval, core.UnshapedArray) for aval in avals), avals
   assert not prim.multiple_results
   weak_type = weak_type_rule(*avals, **kwargs)
@@ -58,8 +59,11 @@ def standard_abstract_eval(prim, shape_rule, dtype_rule, weak_type_rule,
     out = prim.impl(*[x.val for x in avals], **kwargs)
     return core.ConcreteArray(out.dtype, out, weak_type=weak_type)
   elif least_specialized is core.ShapedArray:
+    out_sharding = (sharding_rule(*avals, **kwargs)
+                    if config.sharding_in_types.value else None)
     return core.ShapedArray(shape_rule(*avals, **kwargs),
-                            dtype_rule(*avals, **kwargs), weak_type=weak_type)
+                            dtype_rule(*avals, **kwargs), weak_type=weak_type,
+                            sharding=out_sharding)
   elif least_specialized is core.DShapedArray:
     shape = shape_rule(*avals, **kwargs)
     ty = (core.ShapedArray if all(type(d) is int for d in shape)


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index b4878b021..1e4b5685d 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -4543,6 +4543,32 @@ def spec_regex(s):
   return str(s).replace(r"(", r"\(").replace(r")", r"\)")
 
 
+class ShardingInTypesTest(jtu.JaxTestCase):
+
+  @config.sharding_in_types(True)
+  def test_basic_mul(self):
+    mesh = jtu.create_global_mesh((4, 2), ('x', 'y'))
+    np_inp = np.arange(16).reshape(8, 2)
+    s = NamedSharding(mesh, P('x', 'y'))
+    arr = jax.device_put(np_inp, s)
+
+    @jax.jit
+    def f(x):
+      self.assertEqual(x.sharding.spec, s.spec)
+      x = x * 2
+      self.assertEqual(x.sharding.spec, s.spec)
+      x = x * x
+      self.assertEqual(x.sharding.spec, s.spec)
+      return x
+
+    out = f(arr)
+    self.assertEqual(out.sharding, s)
+    self.assertArraysEqual(out, (np_inp * 2) * (np_inp * 2))
+
+    lowered_text = f.lower(arr).as_text()
+    self.assertEqual(lowered_text.count('@Sharding'), 2)
+
+
 @jtu.pytest_mark_if_available('multiaccelerator')
 class PJitErrorTest(jtu.JaxTestCase):