Merge pull request #197 from ROCm/ci-upstream-sync-83_1

CI: 01/09/25 upstream sync

Author: charleshofer

WORKSPACE

@@ -13,6 +13,7 @@ python_init_repositories(
         "3.11": "//build:requirements_lock_3_11.txt",
         "3.12": "//build:requirements_lock_3_12.txt",
         "3.13": "//build:requirements_lock_3_13.txt",
+        "3.13-ft": "//build:requirements_lock_3_13_ft.txt",
     },
     local_wheel_inclusion_list = [
         "jaxlib*",

build/build.py

@@ -69,7 +69,6 @@ def add_global_arguments(parser: argparse.ArgumentParser):
   parser.add_argument(
       "--python_version",
       type=str,
-      choices=["3.10", "3.11", "3.12", "3.13"],
       default=f"{sys.version_info.major}.{sys.version_info.minor}",
       help=
         """
@@ -390,10 +389,23 @@ async def main():
   bazel_command_base.append("run")
 
   if args.python_version:
+    # Do not add --repo_env=HERMETIC_PYTHON_VERSION with default args.python_version
+    # if bazel_options override it
+    python_version_opt = "--repo_env=HERMETIC_PYTHON_VERSION="
+    if any([python_version_opt in opt for opt in args.bazel_options]):
+      raise RuntimeError(
+        "Please use python_version to set hermetic python version instead of "
+        "setting --repo_env=HERMETIC_PYTHON_VERSION=<python version> bazel option"
+      )
     logging.debug("Hermetic Python version: %s", args.python_version)
     bazel_command_base.append(
         f"--repo_env=HERMETIC_PYTHON_VERSION={args.python_version}"
     )
+    # Let's interpret X.YY-ft version as free-threading python and set rules_python config flag:
+    if args.python_version.endswith("-ft"):
+      bazel_command_base.append(
+        "--@rules_python//python/config_settings:py_freethreaded='yes'"
+      )
 
   # Enable verbose failures.
   bazel_command_base.append("--verbose_failures=true")

build/requirements_lock_3_13_ft.txt

@@ -170,6 +170,7 @@ py_library(
     ] + jax_internal_export_back_compat_test_util_visibility,
     deps = [
         ":jax",
+        ":test_util",
     ] + py_deps("numpy"),
 )
 
@@ -866,6 +867,7 @@ pytype_strict_library(
         ":partition_spec",
         ":sharding",
         ":sharding_specs",
+        ":source_info_util",
         ":tree_util",
         ":util",
         ":xla_bridge",

jax/BUILD

@@ -233,6 +233,8 @@ def trace_context():
 class NoDefault: pass
 no_default = NoDefault()
 
+config_states = {}
+
 class State(config_ext.Config[_T]):
 
   __slots__ = (
@@ -265,6 +267,7 @@ def __init__(
       self._validator(default)
     if self._update_global_hook:
       self._update_global_hook(default)
+    config_states[name] = self
 
   def __bool__(self) -> NoReturn:
     raise TypeError(

jax/_src/config.py

@@ -1696,6 +1696,9 @@ def __init__(self, shape, dtype, weak_type=False, sharding=None):
     self.weak_type = weak_type
     if config.sharding_in_types.value:
       self.sharding = get_sharding(sharding, len(self.shape))
+      if not isinstance(self.sharding.mesh, mesh_lib.AbstractMesh):
+        raise ValueError(
+            f"Mesh of an aval must be an AbstractMesh. Got {self.sharding.mesh}")
 
   def update(self, shape=None, dtype=None, weak_type=None, **kwargs):
     if shape is None:

jax/_src/core.py

@@ -28,13 +28,15 @@
 import jax
 from jax import tree_util
 from jax._src import api_util
+from jax._src import config
 from jax._src import core
 from jax._src import custom_api_util
 from jax._src import dispatch
 from jax._src import linear_util as lu
 from jax._src import mesh as mesh_lib
 from jax._src import sharding_impls
 from jax._src import xla_bridge as xb
+from jax._src.custom_partitioning_sharding_rule import sdy_sharding_rule_to_mlir, SdyShardingRule, str_to_sdy_sharding_rule
 from jax._src.interpreters import mlir
 from jax._src.interpreters import partial_eval as pe
 from jax._src.lib import xla_client as xc
@@ -225,18 +227,20 @@ def _custom_partitioning_abstract_eval(*avals, call, in_tree, out_tree,
                                        propagate_user_sharding, partition,
                                        infer_sharding_from_operands,
                                        decode_shardings,
+                                       sharding_rule,
                                        static_args):
   del in_tree, out_tree, propagate_user_sharding, partition
-  del infer_sharding_from_operands, decode_shardings, static_args
+  del infer_sharding_from_operands, decode_shardings, sharding_rule
+  del static_args
   return call.out_avals
 
 
 def _custom_partitioning_impl(*args, call, in_tree, out_tree,
                               propagate_user_sharding,
                               partition, infer_sharding_from_operands,
-                              decode_shardings, static_args):
+                              decode_shardings, sharding_rule, static_args):
   del in_tree, out_tree, propagate_user_sharding, partition
-  del infer_sharding_from_operands, decode_shardings, static_args
+  del infer_sharding_from_operands, decode_shardings, static_args, sharding_rule
   return core.jaxpr_as_fun(call)(*args)
 
 
@@ -281,7 +285,14 @@ def infer_sharding_from_operands(mesh, arg_shapes, shape):
       arg_shardings = jax.tree.map(lambda x: x.sharding, arg_shapes)
 
 
-    f.def_partition(partition, propagate_user_sharding, infer_sharding_from_operands)
+    f.def_partition(partition, propagate_user_sharding,
+                    infer_sharding_from_operands=infer_sharding_from_operands,
+                    sharding_rule='i j -> 'i j')
+    When config.use_shardy_partitioner.value is True, the sharding_rule is
+    used; otherwise, propagate_user_sharding and infer_sharding_from_operands
+    are used.
+    Instead of using an Einsum-like notation string, sharding_rule can also be
+    a SdyShardingRule object, such as sharding_rule=SdyShardingRule(("i", "j"), ("i", "j")).
 
   The args to ``def_partition`` are as follows:
 
@@ -298,6 +309,10 @@ def infer_sharding_from_operands(mesh, arg_shapes, shape):
   * ``decode_shardings``: When set to True, convert input ``GSPMDSharding``s to
     ``NamedSharding`` if possible. This may not be possible if the user does not
     provide a contextual mesh.
+  * ``sharding_rule``: Either an SdyShardingRule object or an Einsum-like
+    notation string that describes the sharding rule. We borrow the idea from
+    the einops.rearrange string , to use a space separator between factors and
+    allow multiple letters factor names.
 
   Positional arguments can be specified as static using static_argnums. JAX uses
   :code:`inspect.signature(fun)` to resolve these positional arguments.
@@ -350,9 +365,16 @@ def infer_sharding_from_operands(mesh, arg_shapes, result_shape):
       def my_fft(x):
           return fft(x)
 
+      # Use Einsum-like notation to specify the sharding rule.
       my_fft.def_partition(
-          infer_sharding_from_operands=infer_sharding_from_operands,
-          partition=partition)
+        infer_sharding_from_operands=infer_sharding_from_operands,
+        partition=partition,
+        sharding_rule='...i -> ...i')
+      # Use SdyShardingRule object to specify the sharding rule.
+      my_fft.def_partition(
+        infer_sharding_from_operands=infer_sharding_from_operands,
+        partition=partition,
+        sharding_rule=SdyShardingRule(operand_mappings=((SDY_BATCHING, 'i'),), result_mappings=((SDY_BATCHING, 'i'),))))
 
     Now create a 2D array sharded along the first axis, pass it through ``my_fft``
     and notice how it is still sharded as expected, and identical to the output
@@ -425,15 +447,25 @@ def __init__(self, fun, static_argnums=()):
     self.static_argnums = static_argnums
     self.propagate_user_sharding = None
     self.infer_sharding_from_operands = None
+    self.sharding_rule = None
 
   __getattr__: Any = custom_api_util.forward_attr
 
   def def_partition(self, partition, infer_sharding_from_operands,
-                    propagate_user_sharding=None, decode_shardings=True):
+                    propagate_user_sharding=None, decode_shardings=True,
+                    sharding_rule=None):
+    if config.use_shardy_partitioner.value:
+      infer_sharding_from_operands = None
+      propagate_user_sharding = None
+    else:
+      sharding_rule = None
     self.partition = partition
     self.propagate_user_sharding = propagate_user_sharding
     self.infer_sharding_from_operands = infer_sharding_from_operands
     self.decode_shardings = decode_shardings
+    self.sharding_rule = None if sharding_rule is None \
+      else sharding_rule if isinstance(sharding_rule, SdyShardingRule) \
+          else str_to_sdy_sharding_rule(sharding_rule)
     return partition
 
   def __call__(self, *args, **kwargs):
@@ -471,6 +503,7 @@ def __call__(self, *args, **kwargs):
         propagate_user_sharding=self.propagate_user_sharding,
         infer_sharding_from_operands=self.infer_sharding_from_operands,
         decode_shardings=self.decode_shardings,
+        sharding_rule=self.sharding_rule,
         in_tree=in_tree,
         out_tree=out_tree(),
         static_args=static_args
@@ -483,6 +516,7 @@ def _custom_partitioning_lowering_rule(ctx: mlir.LoweringRuleContext, *values,
                                        propagate_user_sharding, partition,
                                        infer_sharding_from_operands,
                                        decode_shardings,
+                                       sharding_rule,
                                        static_args):
   axis_context = ctx.module_context.axis_context
   if (isinstance(axis_context, sharding_impls.SPMDAxisContext) and
@@ -539,6 +573,9 @@ def to_mesh_pspec_sharding(hlo_sharding: xc.HloSharding | None, ndim):
       backend_config=ir.StringAttr.get(key),
       operand_layouts=None,
       result_layouts=None)
+  if sharding_rule is not None:
+    value_types = [mlir.aval_to_ir_type(s) for s in call.in_avals]
+    out.attributes['sdy.sharding_rule'] = sdy_sharding_rule_to_mlir(sharding_rule, value_types, result_types)
   return out.results
 
 mlir.register_lowering(custom_partitioning_p,

jax/_src/custom_partitioning.py

@@ -27,6 +27,9 @@
 # leading ... into factors.
 _BATCHING_DIM_FACTOR_PREFIX = "?"
 
+# A Jax value in general corresponds to an ir.Type or a tuple of ir.Types.
+IrTypes = ir.Type | tuple[ir.Type, ...]
+
 def _check_factor(factor:str):
   """Validates a factor.
 
@@ -278,8 +281,8 @@ def str_to_sdy_sharding_rule(rule: str, **factor_sizes) -> SdyShardingRule:
 
 def sdy_sharding_rule_to_mlir(
   rule: SdyShardingRule,
-  operand_types: list[ir.Type],
-  result_types: list[ir.Type],) -> ir.Attribute:
+  operand_types: list[IrTypes],
+  result_types: list[IrTypes],) -> ir.Attribute:
   """Builds the MLIR representation for the sharding rule.
 
   This is done by verifying that the rule is consistent with the types of
@@ -294,6 +297,10 @@ def sdy_sharding_rule_to_mlir(
     raise ValueError(
       f"Sharding rule has {len(rule.result_mappings)} results, but the operation"
       f" has {len(result_types)} results")
+  if not all(isinstance(t, ir.Type) for t in operand_types + result_types):
+    raise TypeError(
+        f"operand_types and result_types must be a list of ir.Type, but got"
+        f" {operand_types} and {result_types}")
 
   factors_to_indices_sizes: OrderedDict[str, list[int]] = OrderedDict()
   types = operand_types + result_types

jax/_src/custom_partitioning_sharding_rule.py

@@ -49,7 +49,8 @@
 from jax._src.interpreters import xla
 from jax._src.layout import AutoLayout, DeviceLocalLayout
 from jax._src.sharding import Sharding as JSharding
-from jax._src.sharding_impls import AUTO, NamedSharding
+from jax._src.sharding_impls import (AUTO, NamedSharding,
+                                     modify_sdy_sharding_wrt_axis_types)
 from jax._src.lib import xla_client as xc
 from jax._src.lib import xla_extension
 from jax._src.lib.mlir import dialects, ir, passmanager
@@ -1689,13 +1690,17 @@ def lower_jaxpr_to_fun(
           for o, s, o_aval in zip(flat_outputs, ir_result_shardings, output_avals)]
 
     if ir_result_shardings is not None:
-      flat_outputs = [
-          wrap_with_sharding_op(entry_lowering_ctx, o, o_aval, s,
-                                unspecified_dims=us[2])
-          if us[0] and not us[1] else o
-          for o, s, o_aval, us in zip(flat_outputs, ir_result_shardings,
-                                      output_avals, unconstrained_shardings)  # type: ignore
-      ]
+      temp_flat_outputs = []
+      for o, s, o_aval, us in zip(flat_outputs, ir_result_shardings,
+                                  output_avals, unconstrained_shardings):  # type: ignore
+        if us[0] and not us[1]:
+          if config.use_shardy_partitioner.value and config.sharding_in_types.value:
+            s = modify_sdy_sharding_wrt_axis_types(s, o_aval.sharding.mesh)
+          temp_flat_outputs.append(wrap_with_sharding_op(
+              entry_lowering_ctx, o, o_aval, s, unspecified_dims=us[2]))
+        else:
+          temp_flat_outputs.append(o)
+      flat_outputs = temp_flat_outputs
 
     # Insert a custom call if output is on host because XLA needs that to do the
     # transfer.
@@ -2594,14 +2599,20 @@ def lower_sharding_under_shit(ctx, op, aval, sharding_proto=None):
     return op
   # TODO(yashkatariya): If all the axes in pspec are AUTO or collective,
   # `return op` early and avoid bloating HLO size.
-  proto = (aval.sharding._to_xla_hlo_sharding(aval.ndim).to_proto()
-           if sharding_proto is None else sharding_proto)
-  unspecified_dims = None
-  if aval.sharding.mesh._any_axis_collective:
-    unspecified_dims = set(range(aval.ndim))
-  elif aval.sharding.mesh._any_axis_auto:
-    unspecified_dims = {i for i, s in enumerate(aval.sharding.spec) if s is None}
-  return wrap_with_sharding_op(ctx, op, aval, proto, unspecified_dims)
+  if config.use_shardy_partitioner.value:
+    proto = (aval.sharding._to_sdy_sharding(aval.ndim)
+             if sharding_proto is None else sharding_proto)
+    proto = modify_sdy_sharding_wrt_axis_types(proto, aval.sharding.mesh)
+    return wrap_with_sharding_op(ctx, op, aval, proto)
+  else:
+    proto = (aval.sharding._to_xla_hlo_sharding(aval.ndim).to_proto()
+            if sharding_proto is None else sharding_proto)
+    unspecified_dims = None
+    if aval.sharding.mesh._any_axis_auto:
+      # TODO(yashkatariya): Maybe if any mesh axis is auto, mark all axes
+      # as unspecified?
+      unspecified_dims = {i for i, s in enumerate(aval.sharding.spec) if s is None}
+    return wrap_with_sharding_op(ctx, op, aval, proto, unspecified_dims)
 
 
 def set_sharding(op, sharding: xc.OpSharding | sharding_impls.SdyArraySharding):

jax/_src/interpreters/mlir.py

@@ -2163,12 +2163,9 @@ def _abstract_to_concrete_mesh(abstract_mesh):
   out = []
   for s, a in zip(shardings, avals):
     if isinstance(s, UnspecifiedValue) and a.sharding is not None:
-      if config.use_shardy_partitioner.value:
-        spec = a.sharding.spec
-      else:
-        spec = (PartitionSpec(*[PartitionSpec.UNCONSTRAINED if sp is None else sp
-                               for sp in a.sharding.spec])
-                if a.sharding.mesh._any_axis_auto else a.sharding.spec)
+      spec = (PartitionSpec(*[PartitionSpec.UNCONSTRAINED if sp is None else sp
+                              for sp in a.sharding.spec])
+              if a.sharding.mesh._any_axis_auto else a.sharding.spec)
       out.append(NamedSharding(
           _abstract_to_concrete_mesh(a.sharding.mesh), spec))
     else: