Repository: jax-ml/jax
Pull request #27691 (2025-04-03 09:26:47)

[export] Add support for override_lowering_rules to jax.export.
This parameter is already part of the internal API for the AOT lowering function, here we just expose it to `jax.export`.


Patch:
diff --git a/.github/workflows/cloud-tpu-ci-nightly.yml b/.github/workflows/cloud-tpu-ci-nightly.yml
index fd799a3f7..099f4ad5c 100644
--- a/.github/workflows/cloud-tpu-ci-nightly.yml
+++ b/.github/workflows/cloud-tpu-ci-nightly.yml
@@ -26,19 +26,11 @@ jobs:
       matrix:
         jaxlib-version: ["head", "pypi_latest", "nightly", "nightly+oldest_supported_libtpu"]
         tpu: [
+          # {type: "v3-8", cores: "4"}, # Enable when we have the v3 type available
           {type: "v4-8", cores: "4", runner: "linux-x86-ct4p-240-4tpu"},
-          {type: "v5e-8", cores: "8", runner: "linux-x86-ct5lp-224-8tpu"},
-          {type: "v6e-8", cores: "8", runner: "linux-x86-ct6e-180-8tpu"}
+          {type: "v5e-8", cores: "8", runner: "linux-x86-ct5lp-224-8tpu"}
         ]
         python-version: ["3.10"]
-        # Exclude v6e-8 tests for nightly+oldest_supported_libtpu and pypi_latest for resource constraints.
-        exclude:
-          - tpu:
-              type: "v6e-8"
-            jaxlib-version: "nightly+oldest_supported_libtpu"
-          - tpu:
-              type: "v6e-8"
-            jaxlib-version: "pypi_latest"
     name: "TPU test (jaxlib=${{ matrix.jaxlib-version }}, ${{ matrix.tpu.type }})"
     env:
       LIBTPU_OLDEST_VERSION_DATE: 20241205
diff --git a/build/build.py b/build/build.py
index 87aa36aeb..226d984b3 100755
--- a/build/build.py
+++ b/build/build.py
@@ -389,11 +389,6 @@ async def main():
   arch = platform.machine()
   os_name = platform.system().lower()
 
-  custom_wheel_version_suffix = ""
-  wheel_build_date = ""
-  wheel_git_hash = ""
-  wheel_type = "snapshot"
-
   args = parser.parse_args()
 
   logger.info("%s", BANNER)
@@ -626,17 +621,6 @@ async def main():
     )
     for option in args.bazel_options:
       wheel_build_command_base.append(option)
-
-      # Parse the build options for the wheel version suffix.
-      if "ML_WHEEL_TYPE" in option:
-        wheel_type = option.split("=")[-1]
-      if "ML_WHEEL_VERSION_SUFFIX" in option:
-        custom_wheel_version_suffix = option.split("=")[-1].replace("-", "")
-      if "ML_WHEEL_BUILD_DATE" in option:
-        wheel_build_date = option.split("=")[-1].replace("-", "")
-      if "ML_WHEEL_GIT_HASH" in option:
-        wheel_git_hash = option.split("=")[-1][:9]
-
     if "cuda" in args.wheels:
       wheel_build_command_base.append("--config=cuda_libraries_from_stubs")
 
@@ -745,29 +729,10 @@ async def main():
         dst_dir = os.path.join(output_path, wheel_dir)
         utils.copy_dir_recursively(src_dir, dst_dir)
       else:
-        wheel_version_suffix = "dev0+selfbuilt"
-        if wheel_type == "release":
-          wheel_version_suffix = custom_wheel_version_suffix
-        elif wheel_type in ["nightly", "custom"]:
-          wheel_version_suffix = f".dev{wheel_build_date}"
-          if wheel_type == "custom":
-            wheel_version_suffix += (
-                f"+{wheel_git_hash}{custom_wheel_version_suffix}"
-            )
-        if wheel in ["jax", "jax-cuda-pjrt"]:
-          python_tag = "py"
-        else:
-          python_tag = "cp"
-        utils.copy_individual_files(
-            bazel_dir,
-            output_path,
-            f"{wheel_dir}*{wheel_version_suffix}-{python_tag}*.whl",
-        )
+        utils.copy_individual_files(bazel_dir, output_path, f"{wheel_dir}*.whl")
         if wheel == "jax":
           utils.copy_individual_files(
-              bazel_dir,
-              output_path,
-              f"{wheel_dir}*{wheel_version_suffix}.tar.gz",
+              bazel_dir, output_path, f"{wheel_dir}*.tar.gz"
           )
 
   # Exit with success if all wheels in the list were built successfully.
diff --git a/docs/gpu_performance_tips.md b/docs/gpu_performance_tips.md
index b3643cb8e..737486485 100644
--- a/docs/gpu_performance_tips.md
+++ b/docs/gpu_performance_tips.md
@@ -71,10 +71,6 @@ JAX will collect profile information and recompile a module in a single run. Whi
 in manual mode you need to run a task twice, the first time to collect and save profiles
 and the second to compile and run with provided data.
 
-**Important**: the JAX profiler, which is used by both of the PGLE workflows documented
-below, cannot co-exist with the NVIDIA Nsight Systems profiler. This limitation can be
-avoided by using the JAX compilation cache, as described below.
-
 ### Auto PGLE
 The auto PGLE can be turned on by setting the following environment variables:
 
@@ -133,28 +129,6 @@ with config.enable_pgle(True), config.pgle_profiling_runs(1):
   train_step_compiled()
 ```
 
-#### Collecting NVIDIA Nsight Systems profiles when using AutoPGLE
-[jax#24910](https://github.com/jax-ml/jax/pull/24910) (JAX v0.5.1 and newer) added a
-new JAX configuration option, `JAX_COMPILATION_CACHE_EXPECT_PGLE`, which tells JAX to
-attempt to load PGLE-optimized compiled functions from the persistent compilation
-cache.
-
-This allows a two-step process, where the first step writes a PGLE-optimized function
-to the cache:
-```bash
-export JAX_ENABLE_COMPILATION_CACHE=yes          # not strictly needed, on by default
-export JAX_COMPILATION_CACHE_DIR=/root/jax_cache
-JAX_ENABLE_PGLE=yes python my-model.py
-```
-And the second step uses Nsight Systems and loads the PGLE-optimized function from the
-cache:
-```bash
-JAX_COMPILATION_CACHE_EXPECT_PGLE=yes nsys profile python my-model.py
-```
-See also [this page](
-https://docs.jax.dev/en/latest/persistent_compilation_cache.html#pitfalls) for more
-information about the persistent compilation cache and possible pitfalls.
-
 ### Manual PGLE
 
 If you still want to use a manual Profile Guided Latency Estimator the workflow in XLA/GPU is:
diff --git a/jax/_src/config.py b/jax/_src/config.py
index 5b8b87be2..b4a12dcc1 100644
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -998,7 +998,7 @@ explain_cache_misses = bool_state(
     name='jax_explain_cache_misses',
     default=False,
     help=('Each time there is a miss on one of the main caches (e.g. the '
-          'tracing cache), log an explanation.. Logging is performed with '
+          'tracing cache), log an explanation. Logging is performed with '
           '`logging`. When this option is set, the log level is WARNING; '
           'otherwise the level is DEBUG.'))
 
diff --git a/jax/_src/export/_export.py b/jax/_src/export/_export.py
index 9b6a0f809..90cc0c186 100644
--- a/jax/_src/export/_export.py
+++ b/jax/_src/export/_export.py
@@ -529,6 +529,7 @@ def export(
     *,
     platforms: Sequence[str] | None = None,
     disabled_checks: Sequence[DisabledSafetyCheck] = (),
+    _override_lowering_rules: Sequence[tuple[Any, Any]] | None = None
     ) -> Callable[..., Exported]:
   """Exports a JAX function for persistent serialization.
 
@@ -541,6 +542,13 @@ def export(
         If None, then use the default JAX backend.
         The calling convention for multiple platforms is explained at
         https://jax.readthedocs.io/en/latest/export/export.html#module-calling-convention.
+    _override_lowering_rules: an optional sequence of custom lowering rules
+        for some JAX primitives. Each element of the sequence is a pair
+        of a JAX primitive and a lowering function. Defining lowering rules
+        is an advanced feature using JAX internal APIs, which are subject
+        to change. Furthermore, the responsibility for the stability of the
+        MLIR emitted through these custom lowering rules, rests with the user
+        of these rules.
     disabled_checks: the safety checks to disable. See documentation for
         of `jax.export.DisabledSafetyCheck`.
 
@@ -568,7 +576,8 @@ def export(
       Array([0.09983342, 0.19866933, 0.29552022, 0.38941833], dtype=float32)
   """
   return _export_internal(fun_jit, platforms=platforms,
-                          disabled_checks=disabled_checks)
+                          disabled_checks=disabled_checks,
+                          override_lowering_rules=_override_lowering_rules)
 
 
 # TODO(necula): remove this once we improve the integration with jax2tf.
@@ -577,7 +586,8 @@ def _export_internal(
     *,
     platforms: Sequence[str] | None = None,
     disabled_checks: Sequence[DisabledSafetyCheck] = (),
-    _device_assignment_for_internal_jax2tf_use_only = None,
+    _device_assignment_for_internal_jax2tf_use_only=None,
+    override_lowering_rules=None,
     ) -> Callable[..., Exported]:
   """Exports native serialization for a JAX function.
 
@@ -604,6 +614,7 @@ def _export_internal(
     lowered = traced.lower(
         lowering_platforms=actual_lowering_platforms,
         _private_parameters=mlir.LoweringParameters(
+            override_lowering_rules=override_lowering_rules,
             for_export=True,
             export_ignore_forward_compatibility=config.export_ignore_forward_compatibility.value))
     return _export_lowered(
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index a1b37876f..a112063ce 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -663,7 +663,7 @@ class ShapePolyLoweringState:
 @dataclasses.dataclass(frozen=True)
 class LoweringParameters:
   # A mapping between primitives and user-defined LoweringRules.
-  # When lowering a primitive, give priorioty to the rule in this map over
+  # When lowering a primitive, give priority to the rule in this map over
   # existing Jax rules.
   override_lowering_rules: tuple[tuple[core.Primitive, LoweringRule]] | None = None
 
diff --git a/jax/_src/pallas/mosaic_gpu/lowering.py b/jax/_src/pallas/mosaic_gpu/lowering.py
index aafab927d..f027d5bcb 100644
--- a/jax/_src/pallas/mosaic_gpu/lowering.py
+++ b/jax/_src/pallas/mosaic_gpu/lowering.py
@@ -629,9 +629,13 @@ def lower_pipelined_jaxpr_to_module(
     scratch_refs = [
         next(sem_refs_it) if r is sem_placeholder else r for r in scratch_refs
     ]
-    def body_fn(indices, *refs):
+    def body_fn(*refs):
+      grid_env = pallas_core.current_grid_env()
+      assert grid_env is not None  # Set by ``emit_pipeline``.
       program_ids_template = util.merge_lists(
-          which_parallel, indices, [None] * sum(which_parallel)
+          which_parallel,
+          [grid_axis.index for grid_axis in grid_env],
+          [None] * sum(which_parallel),
       )
       assert len(refs) + len(scratch_refs) == len(jaxpr.invars)
       return gpu_primitives.jaxpr_call(
diff --git a/jax/_src/pallas/mosaic_gpu/pipeline.py b/jax/_src/pallas/mosaic_gpu/pipeline.py
index ec088f43c..d85ba4ae2 100644
--- a/jax/_src/pallas/mosaic_gpu/pipeline.py
+++ b/jax/_src/pallas/mosaic_gpu/pipeline.py
@@ -33,6 +33,7 @@ from jax._src.interpreters import partial_eval as pe
 from jax._src.pallas import core as pallas_core
 from jax._src.pallas.mosaic_gpu import core as gpu_core
 from jax._src.pallas.mosaic_gpu import primitives as gpu_primitives
+from jax._src.util import foreach
 from jax.experimental import pallas as pl
 import jax.numpy as jnp
 
@@ -170,8 +171,7 @@ def emit_pipeline(
   """Creates a function to emit a manual pipeline within a Pallas kernel.
 
   Args:
-    body: The pipeline body, called with the indices for the current step, the
-      input refs, followed by the output refs.
+    body: The pipeline body.
     grid: The grid to use for the pipeline.
     in_specs: The block specs for the inputs.
     out_specs: The block specs for the outputs.
@@ -248,8 +248,7 @@ def emit_pipeline(
         it.islice(it.product(*map(range, grid)), max_concurrent_steps)
     ):
       indices = tuple(map(lambda i: jnp.asarray(i, dtype=jnp.int32), indices))
-      for bref in in_brefs:
-        bref.copy_in(step, indices, barrier_ref)
+      foreach(lambda bref: bref.copy_in(step, indices, barrier_ref), in_brefs)
 
     # This is true if any of the outputs need to be transferred inside the loop.
     copies_out_in_loop = not all(bref.is_index_invariant for bref in out_brefs)
@@ -267,13 +266,11 @@ def emit_pipeline(
             max_concurrent_steps - (1 + delay_release), wait_read_only=True
         )
 
-      body(
-          indices,
-          *(
-              bref.get_ref_for_slot(slot)
-              for bref in it.chain(in_brefs, out_brefs)
-          ),
-      )
+      with pallas_core.grid_env(map(pallas_core.GridAxis, indices, grid)):
+        body(*(
+            bref.get_ref_for_slot(slot)
+            for bref in it.chain(in_brefs, out_brefs)
+        ))
 
       if copies_out_in_loop:
         gpu_primitives.commit_smem()
@@ -358,7 +355,6 @@ def emit_pipeline(
 
   return pipeline
 
-
 def emit_pipeline_warp_specialized(
     body: Callable[..., None],
     *,
@@ -380,16 +376,14 @@ def emit_pipeline_warp_specialized(
   ``manual_consumed_barriers`` argument is True.
 
   ```
-  def body(indices, *input_refs, *output_refs, [consumed_barriers]) -> None:
+  def body(*input_refs, *output_refs, [consumed_barriers]) -> None:
   ```
 
   or with a carries enabled (enabled via the ``carry_coroutine`` argument),
   where the body returns the next carry:
 
   ```
-  def body(
-      indices, *input_refs, *output_refs, [consumed_barriers], carry
-  ) -> Carry:
+  def body(*input_refs, *output_refs, [consumed_barriers], carry) -> Carry:
   ```
 
   Args:
@@ -551,17 +545,18 @@ def emit_pipeline_warp_specialized(
         if copies_out_in_loop:
           gpu_primitives.wait_smem_to_gmem(max_concurrent_steps - 1)
 
-        body_refs = []
-        for bref in it.chain(in_brefs, out_brefs):
-          buf_slot = _get_slot(slot, ~bref.is_index_invariant)
-          body_refs.append(bref.get_ref_for_slot(buf_slot))
-
-        body_args = body_refs
-        if manual_consumed_barriers:
-          body_args += [consumed_barrier_ref.at[slot] for consumed_barrier_ref in consumed_barrier_refs]
-        if has_carry:
-          body_args += [prev_body_carry]
-        next_body_carry = body(indices, *body_args)
+        with pallas_core.grid_env(map(pallas_core.GridAxis, indices, grid)):
+          body_refs = []
+          for bref in it.chain(in_brefs, out_brefs):
+            buf_slot = _get_slot(slot, ~bref.is_index_invariant)
+            body_refs.append(bref.get_ref_for_slot(buf_slot))
+
+          body_args = body_refs
+          if manual_consumed_barriers:
+            body_args += [consumed_barrier_ref.at[slot] for consumed_barrier_ref in consumed_barrier_refs]
+          if has_carry:
+            body_args += [prev_body_carry]
+          next_body_carry = body(*body_args)
 
         if not manual_consumed_barriers:
           [consumed_barrier_ref] = consumed_barrier_refs
diff --git a/jax/_src/pallas/mosaic_gpu/primitives.py b/jax/_src/pallas/mosaic_gpu/primitives.py
index 46ec8a870..ff2678454 100644
--- a/jax/_src/pallas/mosaic_gpu/primitives.py
+++ b/jax/_src/pallas/mosaic_gpu/primitives.py
@@ -21,7 +21,7 @@ import dataclasses
 import enum
 import itertools
 import math
-from typing import Any, Literal
+from typing import Any, Literal, Optional
 
 import jax
 from jax._src import core as jax_core
@@ -135,17 +135,17 @@ def _load_p_lowering_rule(
 
 
 def load(
-    src: _Ref, idx, *, layout: Layout | ParameterizedLayout | None = None
-) -> jax.Array:
-  """Loads from a reference into an array with the specified layout.
+    src: _Ref, idx, *, layout: Optional[Layout | ParameterizedLayout] = None
+) -> mgpu.FragmentedArray:
+  """ Loads a ref (SMEM or GMEM) into a FragmentedArray with the specified layout.
 
   Args:
-    src: The reference to load from. Can be either in SMEM or GMEM.
+    src: The reference to copy from.
     idx: The index to load from.
-    layout: The optional layout to use for the resulting array.
+    layout: The optional layout to use for the returned FragmentedArray.
 
   Returns:
-    The loaded array.
+    A FragmentedArray containing the loaded data in the specified layout.
   """
   src, src_transforms = state_primitives.get_ref_and_transforms(
       src, idx, "load", force_trailing_indexer=True,
@@ -160,7 +160,6 @@ def load(
       layout=layout
   )
 
-
 copy_smem_to_gmem_p = jax_core.Primitive("copy_smem_to_gmem")
 copy_smem_to_gmem_p.multiple_results = True
 
@@ -186,7 +185,7 @@ def _copy_smem_to_gmem_lowering(
     dst_transforms_treedef,
     has_user_predicate,
     commit_group,
-    reduction_op,
+    reduction_op: Literal["add", "min", "max", "inc", "dec", "and", "or", "xor"] | None,
 ):
   if has_user_predicate:
     flat_args, user_predicate = flat_args[:-1], flat_args[-1]
@@ -296,7 +295,9 @@ def copy_smem_to_gmem(
     predicate: jax.Array | None = None,
     *,
     commit_group: bool = True,
-    reduction_op: mgpu.ReductionOp | None = None,
+    reduction_op: Literal[
+      "add","min","max","inc","dec","and","or","xor"
+    ] | None = None,
 ) -> None:
   """Asynchronously copies a SMEM reference to a GMEM reference.
 
@@ -305,12 +306,10 @@ def copy_smem_to_gmem(
     dst: The GMEM reference to copy to.
     predicate: A boolean indicating whether the copy should be performed. If
       ``None``, the copy is always performed.
-    commit_group: If ``True``, this and any previously uncommitted copies are
-      committed to a group and can be awaited jointly via
+    commit_group: If ``True``, this and any previously uncommitted copies
+      are committed to a group and can be awaited jointly via
       :func:`jax.experimental.mosaic.gpu.wait_smem_to_gmem`.
-    reduction_op: If set, perform the specified reduction operation when storing
-      to GMEM. For example, using ``"add"`` is conceptually equivalent to
-      doing ``src += dst``.
+    reduction_op: if set, perform the specified reduction op when copy to gmem
 
   See also:
     :func:`jax.experimental.mosaic.gpu.wait_smem_to_gmem`
diff --git a/jax/experimental/mosaic/gpu/__init__.py b/jax/experimental/mosaic/gpu/__init__.py
index e64511594..afc87b5d9 100644
--- a/jax/experimental/mosaic/gpu/__init__.py
+++ b/jax/experimental/mosaic/gpu/__init__.py
@@ -32,7 +32,6 @@ from .core import (
 from .launch_context import (
     LaunchContext as LaunchContext,
     MemRefTransform as MemRefTransform,
-    ReductionOp as ReductionOp,
     Rounding as Rounding,
     TileTransform as TileTransform,
     TransposeTransform as TransposeTransform,
diff --git a/jax/experimental/mosaic/gpu/launch_context.py b/jax/experimental/mosaic/gpu/launch_context.py
index aca3fc723..41c15bc54 100644
--- a/jax/experimental/mosaic/gpu/launch_context.py
+++ b/jax/experimental/mosaic/gpu/launch_context.py
@@ -229,7 +229,6 @@ class CollapseLeadingIndicesTransform(MemRefTransform):
 
 OnDeviceProfiler = profiler.OnDeviceProfiler
 
-ReductionOp = Literal["add", "min", "max", "inc", "dec", "and", "or", "xor"]
 
 @dataclasses.dataclass()
 class LaunchContext:
@@ -407,10 +406,10 @@ class LaunchContext:
       uniform: bool = True,
       collective: Sequence[gpu.Dimension] | gpu.Dimension | None = None,
       partitioned: int | None = None,
-      predicate: (
-          ir.Value | None
-      ) = None,  # Should select 0 or 1 threads from the WG.
-      reduction_op: ReductionOp | None = None,
+      predicate: ir.Value | None = None,  # Should select 0 or 1 threads from the WG.
+      reduction_op: Literal[
+        "add","min","max","inc","dec","and","or","xor"
+      ] | None = None,
   ):
     """Initiates an async copy between GMEM and SMEM.
 
diff --git a/jax/experimental/pallas/ops/gpu/attention_mgpu.py b/jax/experimental/pallas/ops/gpu/attention_mgpu.py
index d06d3b39c..b19e371a1 100644
--- a/jax/experimental/pallas/ops/gpu/attention_mgpu.py
+++ b/jax/experimental/pallas/ops/gpu/attention_mgpu.py
@@ -310,7 +310,7 @@ def attention_with_pipeline_emitter(q, k, v, config: TuningConfig):
       )
       plgpu.wait_smem_to_gmem(0)
 
-    def kv_pipeline(_, k_smem, v_smem,
+    def kv_pipeline(k_smem, v_smem,
                     k_consumed_barrier, v_consumed_barrier,
                     carry):
       acc, m_i, l_i = carry
diff --git a/jaxlib/mlir/_mlir_libs/mosaic_gpu_ext.cc b/jaxlib/mlir/_mlir_libs/mosaic_gpu_ext.cc
index 2751719fc..c73084abc 100644
--- a/jaxlib/mlir/_mlir_libs/mosaic_gpu_ext.cc
+++ b/jaxlib/mlir/_mlir_libs/mosaic_gpu_ext.cc
@@ -138,4 +138,25 @@ NB_MODULE(_mosaic_gpu_ext, m) {
       .def_property_readonly("swizzle", [](MlirAttribute self) {
         return mlirMosaicGpuSwizzleTransformAttrGetSwizzle(self);
       });
+
+  mlir::python::nanobind_adaptors::mlir_attribute_subclass(
+      m, "LayoutAttr", mlirMosaicGpuIsALayoutAttr)
+      .def_classmethod(
+          "get",
+          [](nb::object cls, int32_t num_dimensions,
+             std::vector<MlirAttribute>& transforms, MlirContext ctx) {
+            return cls(mlirMosaicGpuLayoutAttrGet(
+                ctx, num_dimensions, transforms.data(), transforms.size()));
+          },
+          nb::arg("cls"), nb::arg("num_dimensions"), nb::arg("transforms"),
+          nb::arg("context").none() = nb::none(),
+          "Creates a LayoutAttr with the given transforms.")
+      .def_property_readonly("transforms", [](MlirAttribute self) {
+        std::vector<MlirAttribute> result;
+        for (int i = 0; i < mlirMosaicGpuLayoutAttrGetTransformsSize(self);
+             ++i) {
+          result.push_back(mlirMosaicGpuLayoutAttrGetTransform(self, i));
+        }
+        return result;
+      });
 }
diff --git a/jaxlib/mosaic/dialect/gpu/integrations/c/attributes.cc b/jaxlib/mosaic/dialect/gpu/integrations/c/attributes.cc
index 523b14e42..259c37fe5 100644
--- a/jaxlib/mosaic/dialect/gpu/integrations/c/attributes.cc
+++ b/jaxlib/mosaic/dialect/gpu/integrations/c/attributes.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "jaxlib/mosaic/dialect/gpu/integrations/c/attributes.h"
 
 #include <cstdint>
+#include <vector>
 
 #include "mlir-c/IR.h"
 #include "mlir/CAPI/IR.h"
@@ -96,3 +97,36 @@ int32_t mlirMosaicGpuSwizzleTransformAttrGetSwizzle(MlirAttribute attr) {
           .getSwizzle()
           .getValue());
 }
+
+//===----------------------------------------------------------------------===//
+// LayoutAttr
+//===----------------------------------------------------------------------===//
+
+bool mlirMosaicGpuIsALayoutAttr(MlirAttribute attr) {
+  return mlir::isa<mosaic_gpu::LayoutAttr>(unwrap(attr));
+}
+
+MlirAttribute mlirMosaicGpuLayoutAttrGet(MlirContext ctx,
+                                         int32_t num_dimensions,
+                                         MlirAttribute* transforms,
+                                         int32_t transforms_size) {
+  std::vector<mlir::Attribute> unwrapped_transforms;
+  unwrapped_transforms.reserve(transforms_size);
+  for (int i = 0; i < transforms_size; ++i) {
+    unwrapped_transforms.push_back(unwrap(transforms[i]));
+  }
+  return wrap(mosaic_gpu::LayoutAttr::get(unwrap(ctx), num_dimensions,
+                                          unwrapped_transforms));
+}
+
+int32_t mlirMosaicGpuLayoutAttrGetTransformsSize(MlirAttribute attr) {
+  return mlir::cast<mosaic_gpu::LayoutAttr>(unwrap(attr))
+      .getTransforms()
+      .size();
+}
+
+MlirAttribute mlirMosaicGpuLayoutAttrGetTransform(MlirAttribute attr,
+                                                  int32_t index) {
+  return wrap(
+      mlir::cast<mosaic_gpu::LayoutAttr>(unwrap(attr)).getTransforms()[index]);
+}
\ No newline at end of file
diff --git a/jaxlib/mosaic/dialect/gpu/integrations/c/attributes.h b/jaxlib/mosaic/dialect/gpu/integrations/c/attributes.h
index 3221b9220..3b8425b6b 100644
--- a/jaxlib/mosaic/dialect/gpu/integrations/c/attributes.h
+++ b/jaxlib/mosaic/dialect/gpu/integrations/c/attributes.h
@@ -69,6 +69,22 @@ mlirMosaicGpuSwizzleTransformAttrGet(MlirContext ctx, int32_t swizzle);
 MLIR_CAPI_EXPORTED int32_t
 mlirMosaicGpuSwizzleTransformAttrGetSwizzle(MlirAttribute attr);
 
+//===----------------------------------------------------------------------===//
+// LayoutAttr
+//===----------------------------------------------------------------------===//
+
+MLIR_CAPI_EXPORTED bool mlirMosaicGpuIsALayoutAttr(MlirAttribute attr);
+
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirMosaicGpuLayoutAttrGet(MlirContext ctx, int32_t num_dimensions,
+                           MlirAttribute* transforms, int32_t transforms_size);
+
+MLIR_CAPI_EXPORTED int32_t
+mlirMosaicGpuLayoutAttrGetTransformsSize(MlirAttribute attr);
+
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirMosaicGpuLayoutAttrGetTransform(MlirAttribute attr, int32_t index);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/jaxlib/mosaic/dialect/gpu/mosaic_gpu.cc b/jaxlib/mosaic/dialect/gpu/mosaic_gpu.cc
index 073697df5..0a36aab6f 100644
--- a/jaxlib/mosaic/dialect/gpu/mosaic_gpu.cc
+++ b/jaxlib/mosaic/dialect/gpu/mosaic_gpu.cc
@@ -407,6 +407,14 @@ llvm::LogicalResult CustomPrimitiveOp::verify() {
   return llvm::success();
 }
 
+mlir::AffineMap LayoutAttr::getAffineMap() const {
+  // This always returns an identity map. It's technically not correct, but we
+  // don't actually use it anywhere. It's only called during verification of the
+  // layout attribute and needs to be semi-valid.
+  return mlir::AffineMap::getMultiDimIdentityMap(getNumDimensions(),
+                                                 getContext());
+}
+
 void MosaicGPUDialect::initialize() {
   addTypes<
 #define GET_TYPEDEF_LIST
diff --git a/jaxlib/mosaic/dialect/gpu/mosaic_gpu.td b/jaxlib/mosaic/dialect/gpu/mosaic_gpu.td
index cda521855..0d954716b 100644
--- a/jaxlib/mosaic/dialect/gpu/mosaic_gpu.td
+++ b/jaxlib/mosaic/dialect/gpu/mosaic_gpu.td
@@ -225,6 +225,27 @@ def SwizzleTransformAttr : MosaicGPU_Attr<"SwizzleTransform", "swizzle"> {
   let assemblyFormat = "`<` $swizzle `>`";
 }
 
+def LayoutAttr : MosaicGPU_Attr<"Layout", "layout",
+    [DeclareAttrInterfaceMethods<MemRefLayoutAttrInterface>]> {
+  let parameters = (ins
+    TypeParameter<"int32_t", "number of dimensions">:$num_dimensions,
+    ArrayRefParameter<"mlir::Attribute", "transforms">:$transforms
+  );
+
+  let summary = "Specifies a layout of a memref in SMEM.";
+  let description = [{
+    This layout attribute is used to specify the layout of a memref in SMEM.
+    It is composed of a number of transforms, which are applied in the order
+    they are provided. The transforms can be any combination of:
+      - TileTransformAttr
+      - TransposeTransformAttr
+      - SwizzleTransformAttr
+
+    The num_dimensions parameter must match the rank of the memref shape.
+  }];
+  let assemblyFormat = "`<` $num_dimensions `,` $transforms `>`";
+}
+
 def MosaicGPU_AsyncLoadOp : Op<MosaicGPU_Dialect, "async_load",
       [AttrSizedOperandSegments]> {
   let summary = "Schedules an async load of a MemRef from GMEM to SMEM";
@@ -244,9 +265,16 @@ def MosaicGPU_AsyncLoadOp : Op<MosaicGPU_Dialect, "async_load",
     indicates that the slice length is 1 and that the corresponding dimension
     should be collapsed and does not appear in the `destination` MemRef.
 
-    The data is written in row-major order to the contiguous SMEM `destination`.
-    The `source` data does not need to be contiguous, except for the last
-    (and minor-most) dimension.
+    Setting the `layout` attribute of the `destination` MemRef to an instance of
+    `LayoutAttr` enables additional `transforms` that control how the `source`
+    data is mapped to the `destination`. The transformations will be composed in
+    the order they are provided. The transforms may contain up to a single
+    swizzle transform that controls what swizzling is applied to the data after
+    it is transformed, before it is finally written to SMEM. The transformed
+    data is written in row-major order to the contiguous SMEM `destination`.
+    The untransformed `source` data does not need to be contiguous, except for
+    the last dimension, which needs to be contiguous and the minor-most
+    dimension.
 
     The `collective` attribute can be provided to use TMA multicast to more
     efficiently load the GMEM data in cases where multiple thread blocks are
@@ -300,9 +328,14 @@ def MosaicGPU_AsyncStoreOp : Op<MosaicGPU_Dialect, "async_store",
     indicates that this dimension is collapsed in the `source` and needs to be
     expanded to a slice of size 1 in the `destination`.
 
-    The data is written in row-major order to the GMEM `destination`. The
-    `source` data in SMEM needs to be contiguous, but the `destination` GMEM
-    does not.
+    Setting the `layout` attribute of the `source` MemRef to an instance of
+    `LayoutAttr` enables additional `transforms` that control how the
+    `destination` data in GMEM is mapped to the `source` data in SMEM. The
+    transformations will be composed in the order they are provided. The
+    transforms may contain up to a single swizzle transform that is the
+    swizzling mode of the `source` data in SMEM. The `source` SMEM data is
+    contiguous and the transformed data is written to the `destination` GMEM
+    which does not need to be contiguous.
 
     The `predicate` allows scheduling the transfer conditionally. The async copy
     is always scheduled by at most a single lane in the warpgroup.
diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
index c30648a2b..90a19ac95 100644
--- a/third_party/xla/workspace.bzl
+++ b/third_party/xla/workspace.bzl
@@ -21,8 +21,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 #    curl -L https://github.com/openxla/xla/archive/<git hash>.tar.gz | sha256sum
 #    and update XLA_SHA256 with the result.
 
-XLA_COMMIT = "921c164a67e8ac4cf052aab26e849f29b719f802"
-XLA_SHA256 = "9e734da4a0211ac09a00cc07969645e31f107cfee19bbc5d2d1e21ddbb19090d"
+XLA_COMMIT = "c3087e022f3c07f7ed1dd4e47024c437a504341b"
+XLA_SHA256 = "66457303ddec4dbbe43accf38a8b6b635d55808938cf2495443b09ee9c95a147"
 
 def repo():
     tf_http_archive(


Test patch:
diff --git a/.github/workflows/wheel_tests_nightly_release.yml b/.github/workflows/wheel_tests_nightly_release.yml
index 6fd48d016..fd4a52d29 100644
--- a/.github/workflows/wheel_tests_nightly_release.yml
+++ b/.github/workflows/wheel_tests_nightly_release.yml
@@ -80,26 +80,19 @@ jobs:
           tpu-specs: [
             # {type: "v3-8", cores: "4"}, # Enable when we have the v3 type available
             {type: "v4-8", cores: "4", runner: "linux-x86-ct4p-240-4tpu"},
-            {type: "v5e-8", cores: "8", runner: "linux-x86-ct5lp-224-8tpu"},
-            {type: "v6e-8", cores: "8", runner: "linux-x86-ct6e-180-8tpu"}
+            {type: "v5e-8", cores: "8", runner: "linux-x86-ct5lp-224-8tpu"}
           ]
           libtpu-version-type: ["pypi_latest", "nightly", "oldest_supported_libtpu"]
           exclude:
             - libtpu-version-type: ${{ startsWith(github.ref_name, 'release/') && 'nightly' }}
             - libtpu-version-type: ${{ !startsWith(github.ref_name, 'release/') && 'pypi_latest' }}
-            # Run a single Python version for v4-8 and v6e-8
+            # Run a single Python version for v4-8.
             - tpu-specs:
                 type: "v4-8"
               python: "3.10"
             - tpu-specs:
                 type: "v4-8"
               python: "3.11"
-            - tpu-specs:
-                type: "v6e-8"
-              python: "3.10"
-            - tpu-specs:
-                type: "v6e-8"
-              python: "3.11"
             # Run min and max Python versions for v5e-8
             - tpu-specs:
                 type: "v5e-8"
diff --git a/tests/export_test.py b/tests/export_test.py
index 0b78a29a8..2264fbdd9 100644
--- a/tests/export_test.py
+++ b/tests/export_test.py
@@ -19,7 +19,6 @@ import contextlib
 import dataclasses
 import functools
 import logging
-import json
 import math
 import re
 import unittest
@@ -281,6 +280,18 @@ class JaxExportTest(jtu.JaxTestCase):
 
     self.assertAllClose(f(x, y), exp_f.call(x, y))
 
+  def test_override_lowering_rules(self):
+    @jax.jit
+    def f(x):
+      return jnp.sin(x)
+
+    def my_lowering_rule(ctx, arg, **_):
+      return mlir.hlo.CosineOp(arg).results
+
+    exp = get_exported(f, _override_lowering_rules=(
+        (lax.sin_p, my_lowering_rule),))(42.)
+    self.assertIn("stablehlo.cosine", exp.mlir_module())
+
   def test_pytree(self):
     a = np.arange(4, dtype=np.float32)
     b = np.arange(6, dtype=np.float32)
diff --git a/tests/pallas/mosaic_gpu_test.py b/tests/pallas/mosaic_gpu_test.py
index 06dfd453f..acf82ce23 100644
--- a/tests/pallas/mosaic_gpu_test.py
+++ b/tests/pallas/mosaic_gpu_test.py
@@ -1923,7 +1923,7 @@ class PipelineTest(PallasTest):
           max_concurrent_steps=2,
       )(x_gmem, o_gmem)
 
-    def kernel_body(_, x_smem, o_smem):
+    def kernel_body(x_smem, o_smem):
       # +1 for the indexing done by ``emit_pipeline`.
       self.assertLen(x_smem.transforms, len(transforms) + 1)
       o_smem[...] = x_smem[...] + 1.0
@@ -1949,7 +1949,7 @@ class PipelineTest(PallasTest):
           grid=(),
       )(x_gmem, o_gmem)
 
-    def nested_kernel(_, x_gmem, o_gmem):
+    def nested_kernel(x_gmem, o_gmem):
       plgpu.emit_pipeline(
           nested_kernel_body,
           in_specs=[pl.BlockSpec((32, 16), lambda i: (0, i))],
@@ -1958,7 +1958,7 @@ class PipelineTest(PallasTest):
           max_concurrent_steps=2,
       )(x_gmem, o_gmem)
 
-    def nested_kernel_body(_, x_smem, o_smem):
+    def nested_kernel_body(x_smem, o_smem):
       o_smem[...] = x_smem[...] + 1.0
 
     x = jnp.arange(32 * num_steps * 16)
@@ -1983,7 +1983,7 @@ class PipelineTest(PallasTest):
           max_concurrent_steps=2,
       )(x_gmem, o_gmem)
 
-    def kernel_body(_, x_smem, o_smem):
+    def kernel_body(x_smem, o_smem):
       o_smem[...] = x_smem[...] + 1.0
 
     x = jnp.arange(32 * num_steps * 16)
@@ -2016,7 +2016,7 @@ class PipelineTest(PallasTest):
           max_concurrent_steps=2,
       )(x_gmem, o_gmem)
 
-    def kernel_body(_, x_smem, o_smem):
+    def kernel_body(x_smem, o_smem):
       o_smem[...] = x_smem[...] + 1.0
 
     x = jnp.arange(num_steps1 * 32 * num_steps2 * 16)
@@ -2044,7 +2044,7 @@ class PipelineTest(PallasTest):
           max_concurrent_steps=2,
       )(x_gmem, o_gmem)
 
-    def kernel_body(_, x_smem, o_smem):
+    def kernel_body(x_smem, o_smem):
       o_smem[...] = x_smem[...] + 1.0
 
     x = jnp.arange(32 * num_steps1 * 16 * num_steps2 * 8)
@@ -2086,7 +2086,7 @@ class PipelineSm90ATest(PallasSm90ATest):
       )
 
     def kernel(a_gmem, b_gmem, o_smem, acc):
-      def kernel_body(_, a_smem, b_smem):
+      def kernel_body(a_smem, b_smem):
         assert a_smem.shape == (tile_m, tile_k)
         assert b_smem.shape == (tile_k, tile_n)
         plgpu.wgmma(acc, a_smem, b_smem)
@@ -2147,7 +2147,7 @@ class WarpSpecializedPipelineTest(PallasTest):
     x = jax.random.uniform(jax.random.key(0), (m, n), dtype=jnp.float16)
     blk_m = blk_n = 64
 
-    def copy_kernel(_, x_smem, o_smem, o_last_block_smem, *consumed_barriers):
+    def copy_kernel(x_smem, o_smem, o_last_block_smem, *consumed_barriers):
       # TODO(justinfu): Have each wg compute a separate slice
       # after multiple-indexers are supported.
       # This is currently a race, but the values written are the same.
@@ -2201,7 +2201,7 @@ class WarpSpecializedPipelineTest(PallasTest):
         block_shape=(blk_m, blk_n), index_map=lambda i, j: (i, j)
     )
 
-    def tiled_add_kernel(_, x_smem, y_smem, o_smem):
+    def tiled_add_kernel(x_smem, y_smem, o_smem):
       # TODO(justinfu): Have each wg compute a separate slice
       # after multiple-indexers are supported.
       # This is currently a race, but the values written are the same.
@@ -2265,7 +2265,7 @@ class WarpSpecializedPipelineTest(PallasTest):
         plgpu.copy_smem_to_gmem(acc_smem, acc_gmem)
         plgpu.wait_smem_to_gmem(0)
 
-      def tiled_acc_kernel(_, x_smem, carry):
+      def tiled_acc_kernel(x_smem, carry):
         o_carry, = carry
         new_carry = x_smem[...] + o_carry
         return (new_carry,)
@@ -2620,7 +2620,7 @@ class ExamplesTest(PallasTest):
         self.kernel, out_shape=x, grid=(2,), grid_names=("rows",)
     )
     def kernel(l_ref, r_ref, o_ref):
-      def compute(_, l_smem, r_smem, o_smem):
+      def compute(l_smem, r_smem, o_smem):
         o_smem[...] = l_smem[...] + r_smem[...]
       r = lax.axis_index("rows")
       block = pl.BlockSpec((row_block, col_block), lambda c: (r, c))
@@ -2644,7 +2644,7 @@ class ExamplesTest(PallasTest):
         self.kernel, out_shape=x, grid=(2,), grid_names=("rows",)
     )
     def kernel(l_ref, r_ref, o_ref):
-      def compute(_, l_smem, r_smem, o_smem):
+      def compute(l_smem, r_smem, o_smem):
         o_smem[...] = l_smem[...] + r_smem[...]
       r = lax.axis_index("rows")
       block = plgpu.GPUBlockSpec(
@@ -2707,7 +2707,7 @@ class ExamplesSm90ATest(PallasSm90ATest):
         self.kernel, out_shape=x, grid=(2, 2), grid_names=("m", "n")
     )
     def kernel(l_ref, r_ref, o_ref):
-      def compute(_, l_smem, r_smem, o_smem):
+      def compute(l_smem, r_smem, o_smem):
         def do_wgmma(acc_ref):
           plgpu.wgmma(acc_ref, l_smem, r_smem)
           return acc_ref[...]