Merge branch 'main' of https://github.com/jax-ml/jax

Author: nitins17

.github/workflows/ci-build.yaml

@@ -144,13 +144,19 @@ jobs:
 
   documentation_render:
     name: Documentation - render documentation
-    runs-on: ubuntu-latest
+    runs-on: linux-x86-n2-16
+    container:
+      image: index.docker.io/library/ubuntu@sha256:6d8d9799fe6ab3221965efac00b4c34a2bcc102c086a58dff9e19a08b913c7ef # ratchet:ubuntu:20.04
     timeout-minutes: 10
     strategy:
       matrix:
         python-version: ['3.10']
     steps:
     - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+    - name: Image Setup
+      run: |
+        apt update
+        apt install -y libssl-dev libsqlite3-dev
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b  # v5.3.0
       with:
@@ -170,8 +176,7 @@ jobs:
         pip install -r docs/requirements.txt
     - name: Render documentation
       run: |
-        sphinx-build --color -W --keep-going -b html -D nb_execution_mode=off docs docs/build/html
-
+        sphinx-build -j auto --color -W --keep-going -b html -D nb_execution_mode=off docs docs/build/html
 
   jax2tf_test:
     name: "jax2tf_test (py ${{ matrix.python-version }} on ${{ matrix.os }}, x64=${{ matrix.enable-x64}})"

build/build.py

@@ -598,10 +598,12 @@ async def main():
       wheel_build_command.append(f"--jaxlib_git_hash={git_hash}")
 
       result = await executor.run(wheel_build_command.get_command_as_string(), args.dry_run)
+      # Exit with error if any wheel build fails.
       if result.return_code != 0:
         raise RuntimeError(f"Command failed with return code {result.return_code}")
-      else:
-        sys.exit(0)
+
+  # Exit with success if all wheels in the list were built successfully.
+  sys.exit(0)
 
 
 if __name__ == "__main__":

build/rocm/dev_build_rocm.py

@@ -77,13 +77,14 @@ def build_jax_xla(xla_path, rocm_version, rocm_target, use_clang, clang_path):
     build_command = [
         "python3",
         "./build/build.py",
-        "--enable_rocm",
-        "--build_gpu_plugin",
-        "--gpu_plugin_rocm_version=60",
+        "build"
         f"--use_clang={str(use_clang).lower()}",
+        "--wheels=jaxlib,jax-rocm-plugin,jax-rocm-pjrt"
+        "--rocm_path=%/opt/rocm-{rocm_version}/",
+        "--rocm_version=60",
         f"--rocm_amdgpu_targets={rocm_target}",
-        f"--rocm_path=/opt/rocm-{rocm_version}/",
         bazel_options,
+        "--verbose"
     ]
 
     if clang_option:

build/rocm/tools/build_wheels.py

@@ -93,11 +93,12 @@ def build_jaxlib_wheel(
     cmd = [
         "python",
         "build/build.py",
-        "--enable_rocm",
-        "--build_gpu_plugin",
-        "--gpu_plugin_rocm_version=60",
+        "build"
+        "--wheels=jaxlib,jax-rocm-plugin,jax-rocm-pjrt"
         "--rocm_path=%s" % rocm_path,
+        "--rocm_version=60",
         "--use_clang=%s" % use_clang,
+        "--verbose"
     ]
 
     # Add clang path if clang is used.

jax/_src/config.py

@@ -212,6 +212,7 @@ def trace_context():
     return (axis_env_state.value, mesh_context_manager.value,
             xla_metadata_context_manager.value,
             abstract_mesh_context_manager.value,
+            device_context.value,
             compute_on_context_manager.value, enable_x64.value,
             numpy_rank_promotion.value, default_matmul_precision.value,
             dynamic_shapes.value,
@@ -245,6 +246,7 @@ def trace_context():
     axis_env_state = ()
     mesh_context_manager = ()
     abstract_mesh_context_manager = ()
+    device_context = ()
     xla_metadata_context_manager = ()
     compute_on_context_manager = ()
 
@@ -255,12 +257,14 @@ def trace_context():
       mesh_context_manager = context.mesh_context_manager
     if context and context.abstract_mesh_context_manager:
       abstract_mesh_context_manager = context.abstract_mesh_context_manager
+    if context and context.device_context:
+      device_context = context.device_context
     if context and context.xla_metadata_context_manager:
       xla_metadata_context_manager = context.xla_metadata_context_manager
     if context and context.compute_on_context_manager:
       compute_on_context_manager = context.compute_on_context_manager
     return (axis_env_state, mesh_context_manager, abstract_mesh_context_manager,
-            xla_metadata_context_manager,
+            device_context, xla_metadata_context_manager,
             compute_on_context_manager, enable_x64.value,
             numpy_rank_promotion.value, default_matmul_precision.value,
             dynamic_shapes.value,
@@ -976,6 +980,7 @@ def enum_flag(name, default, *args, **kwargs) -> Flag[str]:
   axis_env_state = config_ext.Config((), include_in_jit_key=True)
   mesh_context_manager = config_ext.Config((), include_in_jit_key=True)
   abstract_mesh_context_manager = config_ext.Config((), include_in_jit_key=True)
+  device_context = config_ext.Config((), include_in_jit_key=True)
   compute_on_context_manager = config_ext.Config((), include_in_jit_key=True)
   xla_metadata_context_manager = config_ext.Config((), include_in_jit_key=True)
 else:
@@ -1019,6 +1024,7 @@ class _ThreadLocalExtraJitContext(NamedTuple):
     axis_env_state: Hashable = ()
     mesh_context_manager: Hashable = ()
     abstract_mesh_context_manager: Hashable = ()
+    device_context: Hashable = ()
     compute_on_context_manager: Hashable = ()
     xla_metadata_context_manager: Hashable = ()
 
@@ -1086,6 +1092,7 @@ def set_local(self, value):
   axis_env_state = JitConfig('axis_env_state')
   mesh_context_manager = JitConfig('mesh_context_manager')
   abstract_mesh_context_manager = JitConfig('abstract_mesh_context_manager')
+  device_context = JitConfig('device_context')
   compute_on_context_manager = JitConfig('compute_on_context_manager')
   xla_metadata_context_manager = JitConfig('xla_metadata_context_manager')
 

jax/_src/core.py

@@ -1605,7 +1605,7 @@ def get_sharding(sharding, ndim):
     assert len(sharding.spec) == ndim
     return sharding
 
-  context_mesh = mesh_lib.mesh_context.mesh
+  context_mesh = mesh_lib.abstract_mesh_context.mesh
   # TODO(yashkatariya): Error out and ask users to set the context mesh in their
   # code.
   if context_mesh is None:

jax/_src/interpreters/pxla.py

@@ -2193,8 +2193,15 @@ def lower_sharding_computation(
   assert len(out_shardings) == len(out_layouts) == len(global_out_avals), (
       len(out_shardings), len(out_layouts), len(global_out_avals))
 
-  devices_from_context = (None if context_mesh is None or context_mesh.empty
-                          else context_mesh._flat_devices_tuple)
+  if config.sharding_in_types.value:
+    # TODO(yashkatariya): Thread it via jit path and remove the None check by
+    # making tests go via set_mesh API always.
+    devices_from_context = (
+        None if mesh_lib.device_context.concrete_mesh is None
+        else mesh_lib.device_context.concrete_mesh._flat_devices_tuple)
+  else:
+    devices_from_context = (None if context_mesh is None or context_mesh.empty
+                            else context_mesh._flat_devices_tuple)
   # Device assignment across all inputs, outputs and shardings inside jaxpr
   # should be the same.
   unique_intermediate_shardings = util.stable_unique(

jax/_src/mesh.py

@@ -455,10 +455,10 @@ def local_mesh(self):
     _raise_value_error("local_mesh")
 
   def __enter__(self):
-    return push_mesh_context(self)
+    return push_abstract_mesh_context(self)
 
   def __exit__(self, exc_type, exc_value, traceback):
-    pop_mesh_context()
+    pop_abstract_mesh_context()
     return False
 
   @staticmethod
@@ -473,36 +473,70 @@ def _raise_value_error(name):
   raise ValueError(f"AbstractMesh does not implement {name}")
 
 
-class MeshContext(threading.local):
+class AbstractMeshContext(threading.local):
   def __init__(self):
     self.stack = [None]
     self.mesh = self.stack[-1]
 
-mesh_context = MeshContext()
+abstract_mesh_context = AbstractMeshContext()
 
-def push_mesh_context(val):
-  mesh_context.stack.append(val)
-  mesh_context.mesh = val
+def push_abstract_mesh_context(val):
+  abstract_mesh_context.stack.append(val)
+  abstract_mesh_context.mesh = val
   # TODO(yashkatariya): Allow setting empty tuples and tuples with None in them.
   # Right now that leads to weird numerical issues.
-  non_none_meshes = tuple(m for m in mesh_context.stack if m is not None)
+  non_none_meshes = tuple(m for m in abstract_mesh_context.stack
+                          if m is not None)
   if non_none_meshes:
     jax_config.abstract_mesh_context_manager.set_local(non_none_meshes)
   return val
 
-def pop_mesh_context():
-  mesh_context.stack.pop()
-  mesh_context.mesh = mesh_context.stack[-1]
-  non_none_meshes = tuple(m for m in mesh_context.stack if m is not None)
+def pop_abstract_mesh_context():
+  abstract_mesh_context.stack.pop()
+  abstract_mesh_context.mesh = abstract_mesh_context.stack[-1]
+  non_none_meshes = tuple(m for m in abstract_mesh_context.stack
+                          if m is not None)
   if non_none_meshes:
     jax_config.abstract_mesh_context_manager.set_local(non_none_meshes)
 
 
 class null_mesh_context:
 
   def __enter__(self):
-    return push_mesh_context(None)
+    return push_abstract_mesh_context(None)
 
   def __exit__(self, *excinfo):
-    pop_mesh_context()
+    pop_abstract_mesh_context()
     return False
+
+
+@contextlib.contextmanager
+def set_mesh(mesh: Mesh):
+  with (mesh.abstract_mesh, jax_config.sharding_in_types(True),
+        enter_device_context(mesh)):
+    yield
+
+
+class DeviceContext(threading.local):
+  def __init__(self):
+    self.stack = [None]
+    self.concrete_mesh = self.stack[-1]
+
+device_context = DeviceContext()
+
+
+@contextlib.contextmanager
+def enter_device_context(mesh: Mesh):
+  device_context.stack.append(mesh)
+  device_context.concrete_mesh = mesh
+  non_none_meshes = tuple(m for m in device_context.stack if m is not None)
+  if non_none_meshes:
+    jax_config.device_context.set_local(non_none_meshes)
+  try:
+    yield
+  finally:
+    device_context.stack.pop()
+    device_context.concrete_mesh = device_context.stack[-1]
+    non_none_meshes = tuple(m for m in device_context.stack if m is not None)
+    if non_none_meshes:
+      jax_config.device_context.set_local(non_none_meshes)

jax/_src/pallas/mosaic_gpu/BUILD

@@ -91,7 +91,7 @@ pytype_strict_library(
         ":lowering",
         "//jax",
         "//jax:core",
-        "//jax:effects",
+        "//jax:mlir",
         "//jax:mosaic_gpu",
         "//jax:tree_util",
         "//jax:util",

jax/_src/pallas/mosaic_gpu/primitives.py

@@ -25,8 +25,10 @@
 from jax._src import state
 from jax._src import tree_util
 from jax._src import util
+from jax._src.interpreters import mlir
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import arith as arith_dialect
+from jax._src.lib.mlir.dialects import llvm as llvm_dialect
 from jax._src.lib.mlir.dialects import nvvm as nvvm_dialect
 from jax._src.pallas import core as pallas_core
 from jax._src.pallas.mosaic_gpu import core as gpu_core
@@ -692,3 +694,31 @@ def _commit_smem_lowering(ctx: lowering.LoweringRuleContext):
 def commit_smem():
   """Commits all writes to SMEM, making them visible to loads, TMA and WGMMA."""
   commit_smem_p.bind()
+
+
+broadcasted_iota_p = jax_core.Primitive("broadcasted_iota")
+
+@broadcasted_iota_p.def_abstract_eval
+def _broadcasted_iota_abstract_eval(dtype, shape, dimension, layout):
+  del layout, dimension
+  return jax_core.ShapedArray(shape, dtype)
+
+@lowering.register_lowering_rule(broadcasted_iota_p)
+def _broadcasted_iota_lowering(ctx: lowering.LoweringRuleContext, dtype, shape, dimension, layout):
+  del ctx
+  undef = llvm_dialect.mlir_undef(mlir.dtype_to_ir_type(dtype))
+  is_signed = (
+      jnp.issubdtype(dtype, jnp.signedinteger)
+      if jnp.issubdtype(dtype, jnp.integer)
+      else None
+  )
+  mlir_dtype = mlir.dtype_to_ir_type(dtype)
+  return mgpu.FragmentedArray.splat(
+      undef, shape, layout.value, is_signed=is_signed
+  ).foreach(
+      lambda _, idx: arith_dialect.index_cast(mlir_dtype, idx[dimension]), create_array=True, is_signed=is_signed
+  )
+
+
+def broadcasted_iota(dtype, shape, dimension, *, layout: Layout | None = None):
+  return broadcasted_iota_p.bind(dtype=jnp.dtype(dtype), shape=shape, dimension=dimension, layout=layout)