From ea17f4cb8763f53ab2136f7dec0f79454d917247 Mon Sep 17 00:00:00 2001
From: Naums Mogers <naumsmogers@google.com>
Date: Fri, 18 Sep 2026 10:27:00 -0700
Subject: [PATCH 01/13] [Mosaic] Add helpers to infer sequential accesses from
 tiled layouts

PiperOrigin-RevId: 983936237
---
 .../xla/xla/mosaic/dialect/tpu/util.cc        | 55 +++++++++++++++++++
 third_party/xla/xla/mosaic/dialect/tpu/util.h | 15 +++++
 2 files changed, 70 insertions(+)

diff --git a/third_party/xla/xla/mosaic/dialect/tpu/util.cc b/third_party/xla/xla/mosaic/dialect/tpu/util.cc
index 098d2ba037e873..209de7fbd05a1a 100644
--- a/third_party/xla/xla/mosaic/dialect/tpu/util.cc
+++ b/third_party/xla/xla/mosaic/dialect/tpu/util.cc
@@ -19,13 +19,18 @@ limitations under the License.
 #include <cstdint>
 #include <optional>
 
+#include "absl/log/check.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Location.h"
 #include "mlir/IR/Value.h"
 #include "mlir/IR/ValueRange.h"
 #include "mlir/Support/LLVM.h"
+#include "xla/layout.h"
 #include "xla/mosaic/dialect/tpu/tpu_dialect.h"
 
 namespace mlir::tpu {
@@ -90,4 +95,54 @@ SmallVector<Value> fillPositions(ValueRange values, ArrayRef<int32_t> positions,
   return result;
 }
 
+FailureOr<tpu::TiledLayoutAttr> getTiledLayout(Location loc,
+                                               MemRefType memref_type) {
+  if (auto tiled_layout =
+          dyn_cast<tpu::TiledLayoutAttr>(memref_type.getLayout());
+      tiled_layout != nullptr) {
+    return tiled_layout;
+  }
+  if (memref_type.getLayout().isIdentity()) {
+    return tpu::TiledLayoutAttr::getContiguous(
+        memref_type.getContext(), /*tiles=*/{}, memref_type.getShape());
+  }
+  return emitError(loc, "Not implemented: Unsupported layout");
+}
+
+bool keepsMinorDimSequentialAndContiguous(ArrayRef<int64_t> shape,
+                                          tpu::TiledLayoutAttr layout) {
+  if (shape.empty()) {
+    return true;
+  }
+  const ArrayRef<xla::Tile> tiles = layout.getTiles();
+  const ArrayRef<int64_t> tile_strides = layout.getTileStrides();
+  CHECK(!tile_strides.empty());
+  const bool is_untiled = llvm::all_of(
+      tiles, [](xla::Tile tile) { return tile.dimensions().empty(); });
+  // The memref is untiled and has unit minor stride (or minor shape <= 1).
+  if (is_untiled) {
+    return tile_strides.back() == 1 ||
+           (!ShapedType::isDynamic(shape.back()) && shape.back() <= 1);
+  }
+  if (tiles.size() == 1 && !tiles[0].dimensions().empty()) {
+    const int64_t minor_tile_dim = tiles[0].dimensions().back();
+    const int64_t minor_shape_dim = shape.back();
+    // If the memref does not span more than one tile along the minormost
+    // dimension, the minormost dimension is contiguous within that single tile
+    // regardless of tile strides or other tile dimensions (e.g. 2D tiles).
+    if (!ShapedType::isDynamic(minor_shape_dim) &&
+        minor_shape_dim <= minor_tile_dim) {
+      return true;
+    }
+    // The memref is tiled along the minormost dimension only with consecutive
+    // tiles laid out adjacently (unit minor tile stride).
+    const ArrayRef<int64_t> tile_dims(tiles[0].dimensions().data(),
+                                      tiles[0].dimensions().size());
+    const bool is_single_1d_tile =
+        llvm::all_of(tile_dims.drop_back(), llvm::equal_to<int64_t>(1));
+    return is_single_1d_tile && tile_strides.back() == 1;
+  }
+  return false;
+}
+
 }  // namespace mlir::tpu
diff --git a/third_party/xla/xla/mosaic/dialect/tpu/util.h b/third_party/xla/xla/mosaic/dialect/tpu/util.h
index b2ca879ebbf8a0..18aab9e1b65019 100644
--- a/third_party/xla/xla/mosaic/dialect/tpu/util.h
+++ b/third_party/xla/xla/mosaic/dialect/tpu/util.h
@@ -196,6 +196,21 @@ auto positiveMod(U a, V b) {
 SmallVector<Value> fillPositions(ValueRange values, ArrayRef<int32_t> positions,
                                  int size, Value missing = nullptr);
 
+// Returns the tiled layout of `memref_type`. Memrefs with an identity layout
+// are described by an equivalent untiled layout. Returns an error if the layout
+// is neither identity nor tiled.
+FailureOr<tpu::TiledLayoutAttr> getTiledLayout(Location loc,
+                                               MemRefType memref_type);
+
+// Returns true if the minormost dimension of a memref with the given shape and
+// tiled layout is both sequential (successive elements along that dimension
+// have a physical stride of 1) and contiguous (no intervening padding or
+// non-adjacent runs).
+//
+// This is best-effort and might return false negatives.
+bool keepsMinorDimSequentialAndContiguous(ArrayRef<int64_t> shape,
+                                          tpu::TiledLayoutAttr layout);
+
 }  // namespace mlir::tpu
 
 #endif  // XLA_MOSAIC_DIALECT_TPU_UTIL_H_

From 2e368dbe6a00fc05c082083523e5cac6fafc5474 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Fri, 18 Sep 2026 11:27:11 -0700
Subject: [PATCH 02/13] [XLA:GPU] Fix scatter with empty index vector in
 ScatterWithDistributedIndices

When scattering with an empty index vector (`index_vector_length == 0`) and
sorted/distributed indices:
1. Constructing `thread_id_to_update_id_map` by projecting `indices_map`
   resulted in an empty `IndexingMap` (`IsKnownEmpty()`) because `indices_map`
   has a range variable of size `index_vector_length == 0`, causing the loop
   over updates to have 0 iterations.
2. `EmitInequalityCheck` on empty offset vectors always evaluates to `false`,
   leaving `offsets_changed == false` on the first iteration
   (`iter_slice_id == 0`). Because `is_inbounds_init` starts as `false`,
   `UpdateIsInbounds` never evaluated the bounds check and dropped all updates,
   and the accumulator would combine with 0 instead of being overwritten on the
   first slice.

Fix this by:
- Extracting `ComputeThreadIdToUpdateIdMap` and reusing it to construct
  `indices_map` in `ComputeIndexing`.
- Forcing `offsets_changed` to `true` on the first loop iteration
  (`iter_slice_id == 0`).
- Simplifying loop initialization and update logic:
  - Initialize `indices_init` with `0` instead of `-1` since the first
    iteration check is now explicit.
  - Remove the redundant `arith::SelectOp` loop updating `new_trimmed_offsets`.
  - Remove `is_not_first_iteration` from `write_to_output_required` since
    `iter_is_inbounds` is guaranteed `false` on the first iteration.

PiperOrigin-RevId: 983969974
---
 .../backends/gpu/codegen/emitters/scatter.cc  | 109 ++++++++++--------
 .../backends/gpu/codegen/emitters/scatter.h   |   3 +
 .../tests/scatter/sorted_empty_indices.hlo    |  30 +++++
 third_party/xla/xla/tests/scatter_test.cc     |  32 +++++
 4 files changed, 124 insertions(+), 50 deletions(-)
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_empty_indices.hlo

diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/scatter.cc b/third_party/xla/xla/backends/gpu/codegen/emitters/scatter.cc
index dd96dfbf3b2876..863d024d93bc50 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/scatter.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/scatter.cc
@@ -601,19 +601,14 @@ ScatterWithDistributedIndices::ScatterWithDistributedIndices(
                             num_indices_per_warp_ * num_warps_);
 }
 
-void ScatterWithDistributedIndices::ComputeIndexing(
-    MLIRContext* mlir_context, IndexingMap* updates_map,
-    IndexingMap* indices_map) const {
-  // Compute thread id mapping based on the first update operand.
+IndexingMap ScatterWithDistributedIndices::ComputeThreadIdToUpdateIdMap(
+    MLIRContext* mlir_context) const {
   auto thread_x = CreateDimExpr(MlirKernelFusion::kIndexingMapThreadIdxDims[0],
                                 mlir_context);
   auto block_x = CreateDimExpr(MlirKernelFusion::kIndexingMapBlockIdxDims[0],
                                mlir_context);
   auto warp_id = thread_x / warp_size_;
   auto slice_id = (block_x * num_warps_ + warp_id) / num_warps_per_slice_;
-  auto warp_id_in_slice =
-      (block_x * num_warps_ + warp_id) % num_warps_per_slice_;
-  auto lane_id = thread_x % warp_size_;
   auto index_id_loop = CreateSymbolExpr(0, kGpuGridDims, mlir_context);
   auto index_vector_id = CreateSymbolExpr(1, kGpuGridDims, mlir_context);
 
@@ -623,27 +618,53 @@ void ScatterWithDistributedIndices::ComputeIndexing(
 
   auto grid_vars =
       DimVarsFromGPUGrid({num_warps_ * warp_size_, 1, 1, num_blocks_, 1, 1});
+  IndexingMap thread_id_to_update_id_map{
+      SymbolicMap::Get(mlir_context, kGpuGridDims, /*num_symbols=*/2,
+                       {vectorized_index_id_expr}),
+      grid_vars,
+      {IndexingMap::Variable{
+           {0, num_indices_per_warp_ / indices_vector_size_ - 1},
+           "index_id_loop"},
+       IndexingMap::Variable{{0, indices_vector_size_ - 1}, "index_vector_id"}},
+      /*rt_vars=*/{},
+      {std::make_pair(vectorized_index_id_expr,
+                      Interval{0, description_.num_slices - 1})}};
+  thread_id_to_update_id_map.Simplify();
+  return thread_id_to_update_id_map;
+}
+
+void ScatterWithDistributedIndices::ComputeIndexing(
+    MLIRContext* mlir_context, IndexingMap* updates_map,
+    IndexingMap* indices_map) const {
   if (indices_map) {
+    IndexingMap update_id_map = ComputeThreadIdToUpdateIdMap(mlir_context);
     auto index_dim_loop = CreateSymbolExpr(2, kGpuGridDims, mlir_context);
+    auto range_vars = update_id_map.GetRangeVars();
+    range_vars.push_back(IndexingMap::Variable{
+        {0, description_.index_vector_length - 1}, "index_dim"});
     *indices_map = IndexingMap{
-        SymbolicMap::Get(mlir_context, kGpuGridDims, 3,
-                         {vectorized_index_id_expr, index_dim_loop}),
-        grid_vars,
-        {IndexingMap::Variable{
-             {0, num_indices_per_warp_ / indices_vector_size_ - 1},
-             "index_id_loop"},
-         IndexingMap::Variable{{0, indices_vector_size_ - 1},
-                               "index_vector_id"},
-         IndexingMap::Variable{{0, description_.index_vector_length - 1},
-                               "index_dim"}},
-        /*rt_vars=*/{},
-        {std::make_pair(vectorized_index_id_expr,
-                        Interval{0, description_.num_slices - 1})}};
+        SymbolicMap::Get(
+            mlir_context, kGpuGridDims, /*num_symbols=*/3,
+            {update_id_map.GetSymbolicMap().GetResult(0), index_dim_loop}),
+        update_id_map.GetDimVars(), std::move(range_vars),
+        /*rt_vars=*/{}, update_id_map.GetSymbolicConstraints()};
 
     indices_map->Simplify();
   }
 
   if (updates_map) {
+    // Compute thread id mapping based on the first update operand.
+    auto thread_x = CreateDimExpr(
+        MlirKernelFusion::kIndexingMapThreadIdxDims[0], mlir_context);
+    auto block_x = CreateDimExpr(MlirKernelFusion::kIndexingMapBlockIdxDims[0],
+                                 mlir_context);
+    auto warp_id = thread_x / warp_size_;
+    auto slice_id = (block_x * num_warps_ + warp_id) / num_warps_per_slice_;
+    auto warp_id_in_slice =
+        (block_x * num_warps_ + warp_id) % num_warps_per_slice_;
+    auto lane_id = thread_x % warp_size_;
+    auto grid_vars =
+        DimVarsFromGPUGrid({num_warps_ * warp_size_, 1, 1, num_blocks_, 1, 1});
     auto index_id = CreateSymbolExpr(0, kGpuGridDims, mlir_context);
     auto update_dim_loop = CreateSymbolExpr(1, kGpuGridDims, mlir_context);
     auto vector_id = CreateSymbolExpr(2, kGpuGridDims, mlir_context);
@@ -717,14 +738,8 @@ absl::Status ScatterWithDistributedIndices::EmitEntryFunctionImpl(
   }
   MLIRContext* mlir_context = b.getContext();
 
-  auto thread_id_to_update_id_map = IndexingMap(
-      SymbolicMap::Get(mlir_context, kGpuGridDims, /*num_symbols=*/2,
-                       {indices_map.GetSymbolicMap().GetResult(0)}),
-      indices_map.GetDimVars(),
-      /*range_vars = */
-      {indices_map.GetRangeVars().begin(),
-       indices_map.GetRangeVars().begin() + 2},
-      /*rt vars = */ {}, indices_map.GetSymbolicConstraints());
+  IndexingMap thread_id_to_update_id_map =
+      ComputeThreadIdToUpdateIdMap(mlir_context);
 
   // Convert index_id_loop and index_vector_id to dimension variables.
   IndexingMap slice_indexing =
@@ -735,7 +750,7 @@ absl::Status ScatterWithDistributedIndices::EmitEntryFunctionImpl(
   Value is_inbounds_init = arith::ConstantIntOp::create(b, b.getI1Type(), 0);
   Value slice_id_init = arith::ConstantIndexOp::create(b, 0);
   std::vector<Value> indices_init(description_.index_vector_length,
-                                  arith::ConstantIndexOp::create(b, -1));
+                                  arith::ConstantIndexOp::create(b, 0));
   Value accumulator_init = InitializeAccumulator(b);
   SmallVector<Value> inits =
       Pack({slice_id_init, indices_init, is_inbounds_init, accumulator_init,
@@ -774,32 +789,26 @@ absl::Status ScatterWithDistributedIndices::EmitEntryFunctionImpl(
     auto new_trimmed_offsets =
         helper.ExtractOffsets(nested_b, thread_id_to_index_id_value.front());
 
-    // Check if the offsets changed.
-    Value offsets_changed =
-        EmitInequalityCheck(nested_b, trimmed_offsets, new_trimmed_offsets);
-
-    for (int i = 0; i < description_.index_vector_length; ++i) {
-      new_trimmed_offsets[i] =
-          arith::SelectOp::create(nested_b, offsets_changed,
-                                  new_trimmed_offsets[i], trimmed_offsets[i]);
-    }
+    Value is_first_iteration =
+        arith::CmpIOp::create(nested_b, arith::CmpIPredicate::eq, iter_slice_id,
+                              arith::ConstantIndexOp::create(nested_b, 0));
+    // Check if the offsets changed (or if this is the first iteration).
+    Value offsets_changed = arith::OrIOp::create(
+        nested_b, is_first_iteration,
+        EmitInequalityCheck(nested_b, trimmed_offsets, new_trimmed_offsets));
 
     auto new_offsets = PadWithZeros(new_trimmed_offsets, output_rank,
                                     scatter_indices_to_operand_dims, nested_b);
 
     // Write accumulated values into the tensor if the offsets changed.
-    Value is_not_first_iteration =
-        arith::CmpIOp::create(b, arith::CmpIPredicate::ne, iter_slice_id,
-                              arith::ConstantIndexOp::create(b, 0));
-    Value write_to_output_required = arith::AndIOp::create(
-        b, is_not_first_iteration,
-        arith::AndIOp::create(b, offsets_changed, iter_is_inbounds));
-    iter_output = helper
-                      .WriteAccumulatorsToOutput(b, write_to_output_required,
-                                                 thread_and_block_ids,
-                                                 iter_slice_id, slice_indexing,
-                                                 offsets, iter_acc, iter_output)
-                      .front();
+    Value write_to_output_required =
+        arith::AndIOp::create(nested_b, offsets_changed, iter_is_inbounds);
+    iter_output =
+        helper
+            .WriteAccumulatorsToOutput(
+                nested_b, write_to_output_required, thread_and_block_ids,
+                iter_slice_id, slice_indexing, offsets, iter_acc, iter_output)
+            .front();
 
     // Update `is_inbounds` if the offsets changed.
     Value new_is_inbounds = UpdateIsInbounds(
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/scatter.h b/third_party/xla/xla/backends/gpu/codegen/emitters/scatter.h
index 00a49a8502af88..c48bc2b44cde7d 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/scatter.h
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/scatter.h
@@ -205,6 +205,9 @@ class ScatterWithDistributedIndices : public ScatterFusion {
   // Creates a 2D vector to store the accumulated updates in each thread.
   mlir::Value InitializeAccumulator(mlir::ImplicitLocOpBuilder& b) const;
 
+  IndexingMap ComputeThreadIdToUpdateIdMap(
+      mlir::MLIRContext* mlir_context) const;
+
   // The number of warps that process a single slice of the update.
   int64_t num_warps_per_slice_;
   // The number of indices that every warp iterates over. This is a useful
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_empty_indices.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_empty_indices.hlo
new file mode 100644
index 00000000000000..86cb7efc8f219d
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_empty_indices.hlo
@@ -0,0 +1,30 @@
+// RUN: fusion_to_mlir %s   | emitters_opt -xla-test-optimize \
+// RUN:   -xla-gpu-test-transform-loops  | FileCheck %s
+// RUN: gpu_test_correctness %s --bijection_inputs=scatter:2
+
+add {
+  %p0 = f32[] parameter(0)
+  %p1 = f32[] parameter(1)
+  ROOT %sum = f32[] add(%p0, %p1)
+}
+scatter {
+  %operand = f32[32] parameter(0)
+  %indices = s32[2008,0] parameter(1)
+  %update = f32[2008,32] parameter(2)
+
+  ROOT %scatter = f32[32] scatter(
+      f32[32] %operand,
+      s32[2008,0] %indices,
+      f32[2008,32] %update
+    ),
+    update_window_dims={1},
+    inserted_window_dims={},
+    scatter_dims_to_operand_dims={},
+    index_vector_dim=1,
+    indices_are_sorted=true,
+    unique_indices=false,
+    to_apply=add
+}
+// CHECK-LABEL: func.func @main
+// CHECK: scf.for
+// CHECK:   xla.atomic_rmw
diff --git a/third_party/xla/xla/tests/scatter_test.cc b/third_party/xla/xla/tests/scatter_test.cc
index 664532c03c7934..b8ae5dfe8ade30 100644
--- a/third_party/xla/xla/tests/scatter_test.cc
+++ b/third_party/xla/xla/tests/scatter_test.cc
@@ -1113,6 +1113,38 @@ ENTRY main {
   RunTest(hlo_text, &operand, &scatter_indices, &updates);
 }
 
+TEST_F(ScatterTest, EmptyIndexVectorSortedLarge) {
+  const std::string hlo_text = R"(
+HloModule ScatterEmptyIndexVectorSortedLarge
+
+max_s32 (lhs: s32[], rhs: s32[]) -> s32[] {
+  lhs = s32[] parameter(0)
+  rhs = s32[] parameter(1)
+  ROOT max = s32[] maximum(lhs, rhs)
+}
+
+ENTRY main {
+  operand = s32[1]{0} parameter(0)
+  indices = s32[2000,0]{1,0} parameter(1)
+  updates = s32[2000,1]{1,0} parameter(2)
+  ROOT scatter = s32[1]{0} scatter(operand, indices, updates),
+      update_window_dims={1},
+      inserted_window_dims={},
+      scatter_dims_to_operand_dims={},
+      index_vector_dim=1,
+      indices_are_sorted=true,
+      unique_indices=false,
+      to_apply=max_s32
+}
+)";
+  Literal operand = LiteralUtil::CreateR1<int32_t>({-100});
+  Literal scatter_indices(ShapeUtil::MakeShape(S32, {2000, 0}));
+  Array2D<int32_t> updates_array(2000, 1, -50);
+  updates_array(1234, 0) = -2;
+  Literal updates = LiteralUtil::CreateR2FromArray2D(updates_array);
+  RunTest(hlo_text, &operand, &scatter_indices, &updates);
+}
+
 // Test min/max/add scatters with edge-case values.
 class ScatterEdgeCaseTestP
     : public ScatterTest,

From 639345eb3365f25150cd67dd8c41a5890899731b Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Fri, 18 Sep 2026 12:09:23 -0700
Subject: [PATCH 03/13] Remove exclusive-if-local test tags to eliminate test
 serialization bottlenecks in TensorFlow CUDA CI.

PiperOrigin-RevId: 983993726
---
 tensorflow/python/kernel_tests/distributions/BUILD | 3 ---
 tensorflow/python/kernel_tests/nn_ops/BUILD        | 3 ---
 tensorflow/python/ops/numpy_ops/tests/BUILD        | 1 -
 3 files changed, 7 deletions(-)

diff --git a/tensorflow/python/kernel_tests/distributions/BUILD b/tensorflow/python/kernel_tests/distributions/BUILD
index b37a048bdbb8a5..2e3fa5157f4afd 100644
--- a/tensorflow/python/kernel_tests/distributions/BUILD
+++ b/tensorflow/python/kernel_tests/distributions/BUILD
@@ -326,9 +326,6 @@ cuda_py_strict_test(
         "XLA_PYTHON_CLIENT_PREALLOCATE": "false",
     },
     shard_count = 5,
-    tags = [
-        "exclusive-if-local",
-    ],
     xla_enable_strict_auto_jit = False,
     deps = [
         "//tensorflow/python/eager:backprop",
diff --git a/tensorflow/python/kernel_tests/nn_ops/BUILD b/tensorflow/python/kernel_tests/nn_ops/BUILD
index ab0d0ae3b0511f..dcbf27d841154b 100644
--- a/tensorflow/python/kernel_tests/nn_ops/BUILD
+++ b/tensorflow/python/kernel_tests/nn_ops/BUILD
@@ -909,9 +909,6 @@ cuda_py_strict_test(
         "TF_FORCE_GPU_ALLOW_GROWTH": "true",
         "XLA_PYTHON_CLIENT_PREALLOCATE": "false",
     },
-    tags = [
-        "exclusive-if-local",
-    ],
     deps = [
         ":xent_op_test_base",
         "//tensorflow/python/client:session",
diff --git a/tensorflow/python/ops/numpy_ops/tests/BUILD b/tensorflow/python/ops/numpy_ops/tests/BUILD
index 6542441949ebb7..f7d7e47bd8dc91 100644
--- a/tensorflow/python/ops/numpy_ops/tests/BUILD
+++ b/tensorflow/python/ops/numpy_ops/tests/BUILD
@@ -222,7 +222,6 @@ py_test(
     shard_count = 20,
     strict_deps = True,
     tags = [
-        "exclusive-if-local",
         "gpu",
         "no_pip",
         "requires-gpu-nvidia",

From 89fefb6b26e71095457b4fdde279f5eb261a9760 Mon Sep 17 00:00:00 2001
From: Yang Chen <yangchen@google.com>
Date: Fri, 18 Sep 2026 13:16:35 -0700
Subject: [PATCH 04/13] [tf.data] Propagate context when creating tf.data
 threads.

PiperOrigin-RevId: 984030093
---
 tensorflow/core/kernels/data/BUILD                         | 3 ---
 .../kernels/data/experimental/map_and_batch_dataset_op.cc  | 4 +++-
 .../data/experimental/parallel_interleave_dataset_op.cc    | 7 +++++--
 .../kernels/data/experimental/parse_example_dataset_op.cc  | 7 +++++--
 tensorflow/core/kernels/data/parallel_batch_dataset_op.cc  | 4 +++-
 tensorflow/core/kernels/data/parallel_filter_dataset_op.cc | 4 +++-
 tensorflow/core/kernels/data/parallel_map_dataset_op.cc    | 7 +++++--
 tensorflow/core/kernels/data/prefetch_dataset_op.cc        | 4 +++-
 8 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
index a3df50b0a5e5bf..c3c485580ebdbb 100644
--- a/tensorflow/core/kernels/data/BUILD
+++ b/tensorflow/core/kernels/data/BUILD
@@ -811,7 +811,6 @@ tf_kernel_library(
         "//tensorflow/core:framework_internal",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
-        "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/data:dataset_utils",
         "//tensorflow/core/data:name_utils",
         "//tensorflow/core/data:stats_utils",
@@ -968,11 +967,9 @@ tf_kernel_library(
         "//tensorflow/core:framework_internal",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
-        "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/data:captured_function",
         "//tensorflow/core/data:dataset_utils",
         "//tensorflow/core/data:name_utils",
-        "//tensorflow/core/data:stats_utils",
         "//tensorflow/core/framework:attr_value_proto_cc",
         "//tensorflow/core/framework:types_proto_cc",
         "//tensorflow/core/platform:env",
diff --git a/tensorflow/core/kernels/data/experimental/map_and_batch_dataset_op.cc b/tensorflow/core/kernels/data/experimental/map_and_batch_dataset_op.cc
index 25342ae78db5e7..bdaa80122a8c76 100644
--- a/tensorflow/core/kernels/data/experimental/map_and_batch_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/map_and_batch_dataset_op.cc
@@ -41,6 +41,7 @@ limitations under the License.
 #include "tensorflow/core/platform/stringprintf.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 #include "tensorflow/core/profiler/lib/traceme_encode.h"
+#include "tsl/platform/context.h"
 
 namespace tensorflow {
 namespace data {
@@ -528,7 +529,8 @@ class MapAndBatchDatasetOp::Dataset : public DatasetBase {
         auto new_ctx = std::make_shared<IteratorContext>(*ctx);
         runner_thread_.reset(Env::Default()->StartThread(
             /*thread_options=*/{}, kTFDataMapAndBatch,
-            std::bind(&Iterator::RunnerThread, this, new_ctx)));
+            tsl::WithCurrentContext(
+                std::bind(&Iterator::RunnerThread, this, new_ctx))));
       }
     }
 
diff --git a/tensorflow/core/kernels/data/experimental/parallel_interleave_dataset_op.cc b/tensorflow/core/kernels/data/experimental/parallel_interleave_dataset_op.cc
index 6ec9037254cf13..69f0cac29cb29a 100644
--- a/tensorflow/core/kernels/data/experimental/parallel_interleave_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/parallel_interleave_dataset_op.cc
@@ -36,6 +36,7 @@ limitations under the License.
 #include "tensorflow/core/platform/stringprintf.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 #include "tensorflow/core/profiler/lib/traceme_encode.h"
+#include "tsl/platform/context.h"
 
 namespace tensorflow {
 namespace data {
@@ -574,7 +575,8 @@ class ParallelInterleaveDatasetOp::Dataset : public DatasetBase {
           worker_threads_.emplace_back(Env::Default()->StartThread(
               /*thread_options=*/{},
               absl::StrCat(kDataParallelInterleaveWorker, "_", i),
-              [this, new_ctx, i]() { WorkerThread(new_ctx, i); }));
+              tsl::WithCurrentContext(
+                  [this, new_ctx, i]() { WorkerThread(new_ctx, i); })));
         }
       }
       return absl::OkStatus();
@@ -697,7 +699,8 @@ class ParallelInterleaveDatasetOp::Dataset : public DatasetBase {
           worker_threads_.emplace_back(Env::Default()->StartThread(
               /*thread_options=*/{},
               absl::StrCat(kDataParallelInterleaveWorker, "_", i),
-              [this, new_ctx, i]() { WorkerThread(new_ctx, i); }));
+              tsl::WithCurrentContext(
+                  [this, new_ctx, i]() { WorkerThread(new_ctx, i); })));
         }
         DCHECK(interleave_indices_.size() == dataset()->cycle_length_);
         DCHECK(staging_indices_.size() == dataset()->prefetch_input_elements_);
diff --git a/tensorflow/core/kernels/data/experimental/parse_example_dataset_op.cc b/tensorflow/core/kernels/data/experimental/parse_example_dataset_op.cc
index 018c5f6e6bb0be..81402c4412d9ea 100644
--- a/tensorflow/core/kernels/data/experimental/parse_example_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/parse_example_dataset_op.cc
@@ -37,6 +37,7 @@ limitations under the License.
 #include "tensorflow/core/profiler/lib/traceme.h"
 #include "tensorflow/core/profiler/lib/traceme_encode.h"
 #include "tensorflow/core/util/example_proto_fast_parsing.h"
+#include "tsl/platform/context.h"
 
 namespace tensorflow {
 namespace data {
@@ -571,11 +572,13 @@ class ParseExampleDatasetOp : public UnaryDatasetOpKernel {
           auto ctx_copy = std::make_shared<IteratorContext>(*ctx);
           runner_thread_.reset(Env::Default()->StartThread(
               /*thread_options=*/{}, "tf_data_parallel_map",
-              std::bind(&Iterator::RunnerThread, this, ctx_copy)));
+              tsl::WithCurrentContext(
+                  std::bind(&Iterator::RunnerThread, this, ctx_copy))));
           if (ctx->stats_aggregator()) {
             stats_thread_.reset(Env::Default()->StartThread(
                 /*thread_options=*/{}, "tf_data_parallel_map_stats",
-                std::bind(&Iterator::StatsThread, this, ctx_copy)));
+                tsl::WithCurrentContext(
+                    std::bind(&Iterator::StatsThread, this, ctx_copy))));
           }
         }
       }
diff --git a/tensorflow/core/kernels/data/parallel_batch_dataset_op.cc b/tensorflow/core/kernels/data/parallel_batch_dataset_op.cc
index e70356096509d3..92e22a9c7154ab 100644
--- a/tensorflow/core/kernels/data/parallel_batch_dataset_op.cc
+++ b/tensorflow/core/kernels/data/parallel_batch_dataset_op.cc
@@ -53,6 +53,7 @@ limitations under the License.
 #include "tensorflow/core/profiler/lib/traceme.h"
 #include "tensorflow/core/profiler/lib/traceme_encode.h"
 #include "tensorflow/core/util/batch_util.h"
+#include "tsl/platform/context.h"
 
 namespace tensorflow {
 namespace data {
@@ -480,7 +481,8 @@ class ParallelBatchDatasetOp::Dataset : public DatasetBase {
         auto new_ctx = std::make_shared<IteratorContext>(*ctx);
         runner_thread_.reset(Env::Default()->StartThread(
             /*thread_options=*/{}, kTFDataParallelBatch,
-            std::bind(&Iterator::RunnerThread, this, new_ctx)));
+            tsl::WithCurrentContext(
+                std::bind(&Iterator::RunnerThread, this, new_ctx))));
       }
     }
 
diff --git a/tensorflow/core/kernels/data/parallel_filter_dataset_op.cc b/tensorflow/core/kernels/data/parallel_filter_dataset_op.cc
index cf46bde4242d06..b4b46e78719342 100644
--- a/tensorflow/core/kernels/data/parallel_filter_dataset_op.cc
+++ b/tensorflow/core/kernels/data/parallel_filter_dataset_op.cc
@@ -44,6 +44,7 @@ limitations under the License.
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 #include "tensorflow/core/profiler/lib/traceme_encode.h"
+#include "tsl/platform/context.h"
 
 namespace tensorflow {
 namespace data {
@@ -336,7 +337,8 @@ class ParallelFilterDatasetOp::Dataset : public DatasetBase {
         auto ctx_copy = std::make_shared<IteratorContext>(*ctx);
         runner_thread_.reset(Env::Default()->StartThread(
             /*thread_options=*/{}, "tf_data_parallel_filter",
-            std::bind(&Iterator::RunnerThread, this, ctx_copy)));
+            tsl::WithCurrentContext(
+                std::bind(&Iterator::RunnerThread, this, ctx_copy))));
       }
     }
 
diff --git a/tensorflow/core/kernels/data/parallel_map_dataset_op.cc b/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
index 0870df7f27b115..06d7a89e9ea557 100644
--- a/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
+++ b/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
@@ -56,6 +56,7 @@ limitations under the License.
 #include "tensorflow/core/profiler/lib/traceme.h"
 #include "tensorflow/core/profiler/lib/traceme_encode.h"
 #include "tensorflow/core/protobuf/error_codes.pb.h"
+#include "tsl/platform/context.h"
 
 namespace tensorflow {
 namespace data {
@@ -538,11 +539,13 @@ class ParallelMapDatasetOp::Dataset : public DatasetBase {
         auto ctx_copy = std::make_shared<IteratorContext>(*ctx);
         runner_thread_.reset(Env::Default()->StartThread(
             /*thread_options=*/{}, "tf_data_parallel_map",
-            std::bind(&Iterator::RunnerThread, this, ctx_copy)));
+            tsl::WithCurrentContext(
+                std::bind(&Iterator::RunnerThread, this, ctx_copy))));
         if (ctx->stats_aggregator()) {
           stats_thread_.reset(Env::Default()->StartThread(
               /*thread_options=*/{}, "tf_data_parallel_map_stats",
-              std::bind(&Iterator::StatsThread, this, ctx_copy)));
+              tsl::WithCurrentContext(
+                  std::bind(&Iterator::StatsThread, this, ctx_copy))));
         }
       }
     }
diff --git a/tensorflow/core/kernels/data/prefetch_dataset_op.cc b/tensorflow/core/kernels/data/prefetch_dataset_op.cc
index 0745451e57b0bf..5623c71577572a 100644
--- a/tensorflow/core/kernels/data/prefetch_dataset_op.cc
+++ b/tensorflow/core/kernels/data/prefetch_dataset_op.cc
@@ -53,6 +53,7 @@ limitations under the License.
 #include "tensorflow/core/profiler/lib/traceme.h"
 #include "tensorflow/core/profiler/lib/traceme_encode.h"
 #include "tensorflow/core/protobuf/error_codes.pb.h"
+#include "tsl/platform/context.h"
 #include "tsl/platform/mutex.h"
 
 namespace tensorflow {
@@ -544,7 +545,8 @@ class PrefetchDatasetOp::Dataset : public DatasetBase {
             std::make_shared<IteratorContext>(*ctx);
         prefetch_thread_.reset(Env::Default()->StartThread(
             /*thread_options=*/{}, "tf_data_prefetch",
-            [this, new_ctx]() { PrefetchThread(new_ctx); }));
+            tsl::WithCurrentContext(
+                [this, new_ctx]() { PrefetchThread(new_ctx); })));
       }
       return absl::OkStatus();
     }

From a962b25e8f6fe7a76df8ef38581440938cc016e9 Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Fri, 18 Sep 2026 14:21:39 -0700
Subject: [PATCH 05/13] Increase MSan memory limits and test sharding to
 resolve OOMs and timeouts: - Increase msan_build memory to 32g for
 lazy_sstable_reader_test. - Add msan_build memory (20g) to exec_properties
 for saved_model_test. - Add msan_build memory (30g) to exec_properties for
 remote_test. - Increase shard_count to 16 for linear_operator_block_diag_test
 to prevent MSan timeout.

PiperOrigin-RevId: 984065765
---
 tensorflow/python/distribute/integration_test/BUILD | 1 +
 tensorflow/python/eager/BUILD                       | 6 ++++++
 tensorflow/python/kernel_tests/linalg/BUILD         | 2 +-
 3 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/distribute/integration_test/BUILD b/tensorflow/python/distribute/integration_test/BUILD
index 5a8b3c9aa4837e..00cce0c129efb0 100644
--- a/tensorflow/python/distribute/integration_test/BUILD
+++ b/tensorflow/python/distribute/integration_test/BUILD
@@ -29,6 +29,7 @@ distribute_py_strict_test(
     exec_properties = select({
         # copybara:uncomment_begin(google-only)
         # "@bazel_tools//tools/cpp:asan_build": {"mem": "16g"},
+        # "//tools/cpp:msan_build": {"mem": "20g"},
         # copybara:uncomment_end
         "//conditions:default": None,
     }),
diff --git a/tensorflow/python/eager/BUILD b/tensorflow/python/eager/BUILD
index 2109745801cd1a..18fd0a2b5b71d5 100644
--- a/tensorflow/python/eager/BUILD
+++ b/tensorflow/python/eager/BUILD
@@ -1148,6 +1148,12 @@ cuda_py_strict_test(
     name = "remote_test",
     size = "medium",
     srcs = ["remote_test.py"],
+    exec_properties = select({
+        # copybara:uncomment_begin(google-only)
+        # "//tools/cpp:msan_build": {"mem": "30g"},
+        # copybara:uncomment_end
+        "//conditions:default": None,
+    }),
     grpc_enabled = True,
     shard_count = 2,
     tags = [
diff --git a/tensorflow/python/kernel_tests/linalg/BUILD b/tensorflow/python/kernel_tests/linalg/BUILD
index 39d4158f62a388..6f13559e5dafe7 100644
--- a/tensorflow/python/kernel_tests/linalg/BUILD
+++ b/tensorflow/python/kernel_tests/linalg/BUILD
@@ -202,7 +202,7 @@ cuda_py_strict_test(
     name = "linear_operator_block_diag_test",
     size = "medium",
     srcs = ["linear_operator_block_diag_test.py"],
-    shard_count = 8,
+    shard_count = 16,
     tags = ["optonly"],
     deps = [
         "//tensorflow/python/eager:backprop",

From b53d234eb94b972a392000185cdbc6b8a8369740 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 18 Sep 2026 14:31:06 -0700
Subject: [PATCH 06/13] Miscellaneous fixes for v5p and v6e.

PiperOrigin-RevId: 984070526
---
 .../xla/xla/benchmarks/benchmark_configs.py   | 136 +++++++++++-------
 .../xla/benchmarks/benchmark_configs_test.py  |  72 ++++++----
 .../xla/xla/benchmarks/core/platform_info.py  |   5 +
 .../pallas_microbenchmarks/cost_model.py      |  13 +-
 .../pallas_microbenchmarks/dense_matmul.py    |   7 +-
 .../dense_matmul_lib.py                       |   9 +-
 .../subchannel_matmul.py                      |  14 +-
 .../subchannel_matmul_lib.py                  |  23 +--
 8 files changed, 182 insertions(+), 97 deletions(-)

diff --git a/third_party/xla/xla/benchmarks/benchmark_configs.py b/third_party/xla/xla/benchmarks/benchmark_configs.py
index c360569c19e6fc..d13b0155b4acdf 100644
--- a/third_party/xla/xla/benchmarks/benchmark_configs.py
+++ b/third_party/xla/xla/benchmarks/benchmark_configs.py
@@ -21,6 +21,7 @@
 from jax.experimental.pallas import tpu as pltpu
 import jax.numpy as jnp
 
+from xla.benchmarks.core import platform_info
 from xla.benchmarks.dma_microbenchmarks import chip_to_chip_dma_benchmark
 from xla.benchmarks.dma_microbenchmarks import chiplet_to_chiplet_dma_benchmark
 from xla.benchmarks.dma_microbenchmarks import host_dma_benchmark
@@ -33,13 +34,7 @@
 
 _DIM_VALUES = (1024, 2048, 4096, 8192, 16384, 32768)
 
-_LHS_RHS_DTYPE_PAIRS = (
-    (jnp.bfloat16, jnp.bfloat16),
-    (jnp.bfloat16, jnp.float8_e4m3fn),
-    (jnp.bfloat16, jnp.int4),
-    (jnp.float8_e4m3fn, jnp.float8_e4m3fn),
-    (jnp.float8_e4m3fn, jnp.int4),
-)
+_SUBCHANNEL_DIM_VALUES = (4096, 8192, 16384)
 
 _OUT_DTYPE_PAIRS = (
     jnp.float32,
@@ -47,18 +42,38 @@
 )
 
 
+def _lhs_rhs_dtype_pairs(
+    low_precision_dtype: jnp.dtype,
+) -> list[tuple[jnp.dtype, jnp.dtype]]:
+  return [  # pyrefly: ignore[bad-return]
+      (jnp.bfloat16, jnp.bfloat16),
+      (jnp.bfloat16, low_precision_dtype),
+      (low_precision_dtype, low_precision_dtype),
+  ]
+
+
 def get_dense_matmul_configs(
     chip_version: pltpu.ChipVersion | None = None,
 ) -> list[dense_matmul_lib.DenseMatmulConfig]:
   """Generates preconfigured dense matmul benchmark configs."""
   configs = []
-  acc_dtype = jnp.float32
+  pinfo = platform_info.get_platform_info(chip_version)
+  if jnp.int4 in pinfo.matmul_cadence_cycles_by_dtype:
+    low_precision_dtype = jnp.int4
+  else:
+    low_precision_dtype = jnp.float8_e4m3fn
   subblock_m = dense_matmul_lib.get_default_subblock_m(chip_version)
 
   for m in _DIM_VALUES:
     n = k = m
     mem_options = [pltpu.HBM, pltpu.VMEM] if m in (1024, 2048) else [pltpu.HBM]
-    for lhs_dtype, rhs_dtype in _LHS_RHS_DTYPE_PAIRS:
+    for lhs_dtype, rhs_dtype in _lhs_rhs_dtype_pairs(
+        low_precision_dtype,  # pyrefly: ignore[bad-argument-type]
+    ):
+      if jnp.issubdtype(lhs_dtype, jnp.integer) and pinfo.generation < 8:
+        acc_dtype = jnp.int32
+      else:
+        acc_dtype = jnp.float32
       for out_dtype in _OUT_DTYPE_PAIRS:
         for mem in mem_options:
           block_m, block_k, block_n = dense_matmul_lib.select_window(
@@ -101,52 +116,61 @@ def get_subchannel_matmul_configs(
 ) -> list[subchannel_matmul_lib.SubchannelMatmulConfig]:
   """Generates preconfigured subchannel matmul benchmark configs."""
   configs = []
-  m, k, n = 128, 8192, 4096
+  pinfo = platform_info.get_platform_info(chip_version)
+  m = 128
   lhs_dtype = rhs_dtype = out_dtype = jnp.bfloat16
-  acc_dtype = jnp.bfloat16
+  outer_acc_dtype = jnp.bfloat16
   subchannel_size = 1024
-  lhs_quantized_dtype = jnp.float8_e4m3fn
+  if jnp.int4 in pinfo.matmul_cadence_cycles_by_dtype:
+    lhs_quantized_dtype = jnp.int4
+    inner_acc_dtype = jnp.int32
+  else:
+    lhs_quantized_dtype = jnp.float8_e4m3fn
+    inner_acc_dtype = jnp.float32
   rhs_quantized_dtype = jnp.int4
   pre_quantize_lhs = False
-
-  for mem in [pltpu.HBM, pltpu.VMEM]:
-    block_m, block_k, block_n = subchannel_matmul_lib.select_window(
-        m=m,
-        k=k,
-        n=n,
-        lhs_mem=mem,
-        rhs_mem=mem,
-        out_mem=mem,
-        lhs_dtype=lhs_dtype,
-        rhs_dtype=rhs_dtype,
-        out_dtype=out_dtype,
-        acc_dtype=acc_dtype,
-        lhs_quantized_dtype=lhs_quantized_dtype,
-        rhs_quantized_dtype=rhs_quantized_dtype,
-        pre_quantize_lhs=pre_quantize_lhs,
-        chip_version=chip_version,
-    )
-    configs.append(
-        subchannel_matmul_lib.SubchannelMatmulConfig(
-            m=m,
-            k=k,
-            n=n,
-            block_m=int(block_m),
-            block_k=int(block_k),
-            block_n=int(block_n),
-            subchannel_size=subchannel_size,
-            lhs_mem=mem,
-            rhs_mem=mem,
-            out_mem=mem,
-            lhs_dtype=lhs_dtype,
-            rhs_dtype=rhs_dtype,
-            out_dtype=out_dtype,
-            acc_dtype=acc_dtype,
-            lhs_quantized_dtype=lhs_quantized_dtype,
-            rhs_quantized_dtype=rhs_quantized_dtype,
-            pre_quantize_lhs=pre_quantize_lhs,
-        )
-    )
+  for k in _SUBCHANNEL_DIM_VALUES:
+    n = k
+    mem_options = [pltpu.HBM, pltpu.VMEM] if k == 4096 else [pltpu.HBM]
+    for mem in mem_options:
+      block_m, block_k, block_n = subchannel_matmul_lib.select_window(
+          m=m,
+          k=k,
+          n=n,
+          lhs_mem=mem,
+          rhs_mem=mem,
+          out_mem=mem,
+          lhs_dtype=lhs_dtype,
+          rhs_dtype=rhs_dtype,
+          out_dtype=out_dtype,
+          outer_acc_dtype=outer_acc_dtype,
+          lhs_quantized_dtype=lhs_quantized_dtype,
+          rhs_quantized_dtype=rhs_quantized_dtype,
+          pre_quantize_lhs=pre_quantize_lhs,
+          chip_version=chip_version,
+      )
+      configs.append(
+          subchannel_matmul_lib.SubchannelMatmulConfig(
+              m=m,
+              k=k,
+              n=n,
+              block_m=int(block_m),
+              block_k=int(block_k),
+              block_n=int(block_n),
+              subchannel_size=subchannel_size,
+              lhs_mem=mem,
+              rhs_mem=mem,
+              out_mem=mem,
+              lhs_dtype=lhs_dtype,
+              rhs_dtype=rhs_dtype,
+              out_dtype=out_dtype,
+              inner_acc_dtype=inner_acc_dtype,
+              outer_acc_dtype=outer_acc_dtype,
+              lhs_quantized_dtype=lhs_quantized_dtype,
+              rhs_quantized_dtype=rhs_quantized_dtype,
+              pre_quantize_lhs=pre_quantize_lhs,
+          )
+      )
   return configs
 
 
@@ -154,11 +178,17 @@ def get_jax_matmul_configs(
     chip_version: pltpu.ChipVersion | None = None,
 ) -> list[matmul_lib.JaxMatmulConfig]:
   """Generates preconfigured JAX matmul benchmark configs."""
-  del chip_version  # Unused.
   configs = []
+  pinfo = platform_info.get_platform_info(chip_version)
+  if jnp.int4 in pinfo.matmul_cadence_cycles_by_dtype:
+    low_precision_dtype = jnp.int4
+  else:
+    low_precision_dtype = jnp.float8_e4m3fn
   for m in _DIM_VALUES:
     n = k = m
-    for lhs_dtype, rhs_dtype in _LHS_RHS_DTYPE_PAIRS:
+    for lhs_dtype, rhs_dtype in _lhs_rhs_dtype_pairs(
+        low_precision_dtype,  # pyrefly: ignore[bad-argument-type]
+    ):
       for out_dtype in _OUT_DTYPE_PAIRS:
         configs.append(
             matmul_lib.JaxMatmulConfig(
diff --git a/third_party/xla/xla/benchmarks/benchmark_configs_test.py b/third_party/xla/xla/benchmarks/benchmark_configs_test.py
index 75b4636e84c822..acfa18f988de10 100644
--- a/third_party/xla/xla/benchmarks/benchmark_configs_test.py
+++ b/third_party/xla/xla/benchmarks/benchmark_configs_test.py
@@ -31,16 +31,22 @@
 
 class BenchmarkConfigsTest(parameterized.TestCase):
 
-  def test_dense_matmul_configs(self):
-    chip = pltpu.ChipVersion.TPU_V5E
+  @parameterized.named_parameters(
+      ("v5e", pltpu.ChipVersion.TPU_V5E, jnp.int4, jnp.int32),
+      ("7x", pltpu.ChipVersion.TPU_7X, jnp.float8_e4m3fn, jnp.float32),
+      ("8t", pltpu.ChipVersion.TPU_8T, jnp.int4, jnp.float32),
+  )
+  def test_dense_matmul_configs(
+      self, chip, expected_low_precision_dtype, expected_int_acc_dtype
+  ):
     configs = benchmark_configs.get_dense_matmul_configs(chip_version=chip)
     # Expected total configs:
     # M in [1024, 2048] ->
-    #   2 sizes * 5 dtype pairs * 2 out_dtypes * 2 mems (HBM, VMEM) = 40
+    #   2 sizes * 3 dtype pairs * 2 out_dtypes * 2 mems (HBM, VMEM) = 24
     # M in [4096, 8192, 16384, 32768] ->
-    #   4 sizes * 4 dtype pairs * 2 out_dtypes * 1 mem (HBM) = 40
-    # Total = 32 configs.
-    self.assertLen(configs, 80)
+    #   4 sizes * 3 dtype pairs * 2 out_dtypes * 1 mem (HBM) = 24
+    # Total = 48 configs.
+    self.assertLen(configs, 48)
 
     for cfg in configs:
       self.assertIsInstance(cfg, dense_matmul_lib.DenseMatmulConfig)
@@ -51,14 +57,15 @@ def test_dense_matmul_configs(self):
           (cfg.lhs_dtype, cfg.rhs_dtype),
           [
               (jnp.bfloat16, jnp.bfloat16),
-              (jnp.bfloat16, jnp.float8_e4m3fn),
-              (jnp.bfloat16, jnp.int4),
-              (jnp.float8_e4m3fn, jnp.float8_e4m3fn),
-              (jnp.float8_e4m3fn, jnp.int4),
+              (jnp.bfloat16, expected_low_precision_dtype),
+              (expected_low_precision_dtype, expected_low_precision_dtype),
           ],
       )
       self.assertIn(cfg.out_dtype, [jnp.float32, jnp.bfloat16])
-      self.assertEqual(cfg.acc_dtype, jnp.float32)
+      if jnp.issubdtype(cfg.lhs_dtype, jnp.integer):
+        self.assertEqual(cfg.acc_dtype, expected_int_acc_dtype)
+      else:
+        self.assertEqual(cfg.acc_dtype, jnp.float32)
       self.assertEqual(cfg.lhs_mem, cfg.rhs_mem)
       self.assertEqual(cfg.lhs_mem, cfg.out_mem)
 
@@ -71,27 +78,38 @@ def test_dense_matmul_configs(self):
       self.assertGreater(cfg.block_k, 0)
       self.assertGreater(cfg.block_n, 0)
 
-  def test_subchannel_matmul_configs(self):
-    chip = pltpu.ChipVersion.TPU_V5E
+  @parameterized.named_parameters(
+      ("v5e", pltpu.ChipVersion.TPU_V5E, jnp.int4, jnp.int32),
+      ("7x", pltpu.ChipVersion.TPU_7X, jnp.float8_e4m3fn, jnp.float32),
+  )
+  def test_subchannel_matmul_configs(
+      self, chip, expected_lhs_quantized_dtype, expected_inner_acc_dtype
+  ):
     configs = benchmark_configs.get_subchannel_matmul_configs(chip_version=chip)
-    # Expected: 2 configs (HBM and VMEM)
-    self.assertLen(configs, 2)
+    # Expected: 4 configs (k=4096 with HBM and VMEM; k=8192, 16384 with HBM)
+    self.assertLen(configs, 4)
 
     mems = set()
     for cfg in configs:
       self.assertIsInstance(cfg, subchannel_matmul_lib.SubchannelMatmulConfig)
       self.assertEqual(cfg.m, 128)
-      self.assertEqual(cfg.k, 8192)
-      self.assertEqual(cfg.n, 4096)
+      self.assertEqual(cfg.k, cfg.n)
+      self.assertIn(cfg.k, [4096, 8192, 16384])
       self.assertEqual(cfg.lhs_dtype, jnp.bfloat16)
       self.assertEqual(cfg.rhs_dtype, jnp.bfloat16)
       self.assertEqual(cfg.out_dtype, jnp.bfloat16)
+      self.assertEqual(cfg.inner_acc_dtype, expected_inner_acc_dtype)
+      self.assertEqual(cfg.outer_acc_dtype, jnp.bfloat16)
       self.assertEqual(cfg.subchannel_size, 1024)
-      self.assertEqual(cfg.lhs_quantized_dtype, jnp.float8_e4m3fn)
+      self.assertEqual(cfg.lhs_quantized_dtype, expected_lhs_quantized_dtype)
       self.assertEqual(cfg.rhs_quantized_dtype, jnp.int4)
       self.assertFalse(cfg.pre_quantize_lhs)
       self.assertEqual(cfg.lhs_mem, cfg.rhs_mem)
       self.assertEqual(cfg.lhs_mem, cfg.out_mem)
+      if cfg.k == 4096:
+        self.assertIn(cfg.lhs_mem, [pltpu.HBM, pltpu.VMEM])
+      else:
+        self.assertEqual(cfg.lhs_mem, pltpu.HBM)
       self.assertGreater(cfg.block_m, 0)
       self.assertGreater(cfg.block_k, 0)
       self.assertGreater(cfg.block_n, 0)
@@ -99,10 +117,14 @@ def test_subchannel_matmul_configs(self):
 
     self.assertEqual(mems, {pltpu.HBM, pltpu.VMEM})
 
-  def test_jax_matmul_configs(self):
-    configs = benchmark_configs.get_jax_matmul_configs()
-    # 6 dim sizes * 5 dtype pairs * 2 out dtypes = 60 configs.
-    self.assertLen(configs, 60)
+  @parameterized.named_parameters(
+      ("v5e", pltpu.ChipVersion.TPU_V5E, jnp.int4),
+      ("7x", pltpu.ChipVersion.TPU_7X, jnp.float8_e4m3fn),
+  )
+  def test_jax_matmul_configs(self, chip, expected_low_precision_dtype):
+    configs = benchmark_configs.get_jax_matmul_configs(chip_version=chip)
+    # 6 dim sizes * 3 dtype pairs * 2 out dtypes = 36 configs.
+    self.assertLen(configs, 36)
     for cfg in configs:
       self.assertIsInstance(cfg, matmul_lib.JaxMatmulConfig)
       self.assertEqual(cfg.b, 1)
@@ -113,10 +135,8 @@ def test_jax_matmul_configs(self):
           (cfg.lhs_dtype, cfg.rhs_dtype),
           [
               (jnp.bfloat16, jnp.bfloat16),
-              (jnp.bfloat16, jnp.float8_e4m3fn),
-              (jnp.bfloat16, jnp.int4),
-              (jnp.float8_e4m3fn, jnp.float8_e4m3fn),
-              (jnp.float8_e4m3fn, jnp.int4),
+              (jnp.bfloat16, expected_low_precision_dtype),
+              (expected_low_precision_dtype, expected_low_precision_dtype),
           ],
       )
       self.assertIn(cfg.out_dtype, [jnp.float32, jnp.bfloat16])
diff --git a/third_party/xla/xla/benchmarks/core/platform_info.py b/third_party/xla/xla/benchmarks/core/platform_info.py
index 2a8d643616cb16..af6f0c562ec742 100644
--- a/third_party/xla/xla/benchmarks/core/platform_info.py
+++ b/third_party/xla/xla/benchmarks/core/platform_info.py
@@ -97,6 +97,11 @@ def num_accumulators(self) -> int:
     """Returns the number of accumulators on the TPU platform."""
     return self._tpu_info.num_accumulators
 
+  @property
+  def generation(self) -> int:
+    """Returns the generation of the TPU platform."""
+    return self._tpu_info.generation
+
   def get_emulated_dtype(
       self, dtype: jnp.dtype, supported_dtypes: Mapping[jnp.dtype, int]
   ) -> jnp.dtype:
diff --git a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/cost_model.py b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/cost_model.py
index ca8ec7db9a9bc4..01af02c1b86cc8 100644
--- a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/cost_model.py
+++ b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/cost_model.py
@@ -349,21 +349,28 @@ def __init__(
         sp_n,
         sp_m,
     )
-    # If the LHS is emulated as a different dtype, then the compiler may spill
+    # If the LHS requires any VPU processing, then the compiler may spill
     # the converted LHS operand to VMEM, so we account for this below.
+    spill_bits = None
     emulated_lhs = (
         lhs_dtype not in self._platform_info.matmul_cadence_cycles_by_dtype
     )
     if emulated_lhs:
+      # Emulated dtypes need VPU for unpacking and conversion.
       emulated_dtype = self._platform_info.get_emulated_dtype(
           lhs_dtype,
           self._platform_info.matmul_cadence_cycles_by_dtype,
       )
-      emulated_bits = jax.dtypes.itemsize_bits(emulated_dtype)
+      spill_bits = jax.dtypes.itemsize_bits(emulated_dtype)
+    elif lhs_dtype == jnp.int4 and self._platform_info.generation < 8:
+      # Prior to generation 8, int4 LHS requires VPU for unpacking from
+      # compressed format and repacking to interleaved format.
+      spill_bits = 4
+    if spill_bits is not None:
       block_m = self._block_m
       if subblock_m is not None:
         block_m = np.minimum(self._block_m, subblock_m)
-      self._vmem_usage_bytes += block_m * self._block_k * emulated_bits // 8
+      self._vmem_usage_bytes += block_m * self._block_k * spill_bits // 8
 
   def _compute_matmul_latency_ns_per_iteration(
       self, p_state: int | None
diff --git a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul.py b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul.py
index 529a6e47f02de1..23b89a0ad1d3e2 100644
--- a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul.py
+++ b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul.py
@@ -21,6 +21,7 @@
 import jax.numpy as jnp
 
 from xla.benchmarks.core import benchmark
+from xla.benchmarks.core import platform_info
 from xla.benchmarks.pallas_microbenchmarks import dense_matmul_lib
 
 _DIM = flags.DEFINE_list(
@@ -94,7 +95,11 @@ def main(_):
   lhs_dtype = benchmark.str_to_dtype(_FMT.value[0])
   rhs_dtype = benchmark.str_to_dtype(_FMT.value[1])
   out_dtype = benchmark.str_to_dtype(_FMT.value[2])
-  acc_dtype = jnp.float32
+  pinfo = platform_info.get_platform_info()
+  if jnp.issubdtype(lhs_dtype, jnp.integer) and pinfo.generation < 8:
+    acc_dtype = jnp.int32
+  else:
+    acc_dtype = jnp.float32
 
   if len(_DIM.value) != 4:
     raise ValueError(f"Expected 4 dims, got {len(_DIM.value)}")
diff --git a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py
index 7b4cd08420df62..44a34f79ad6d6b 100644
--- a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py
+++ b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py
@@ -36,6 +36,8 @@
 
 _KERNEL_NAME_TEMPLATE = "matmul_{m}_{k}_{n}_{lhs_dtype}_{rhs_dtype}_{out_dtype}"
 
+_NO_ACCUMULATORS_DEFAULT_SUBBLOCK_M = 256
+
 
 def get_default_subblock_m(
     chip_version: pltpu.ChipVersion | None = None,
@@ -46,7 +48,12 @@ def get_default_subblock_m(
   if num_accumulators > 0:
     num_sublanes, _ = pinfo.vreg_size
     return num_sublanes * num_accumulators
-  return None
+  else:
+    # Accumulation occurs in vregs + vmem, so there may be additional spills
+    # that the cost model doesn't account for if there are too many
+    # accumulations in-flight. So we set subblock_m to a relatively low value
+    # to avoid this.
+    return _NO_ACCUMULATORS_DEFAULT_SUBBLOCK_M
 
 
 def select_window(
diff --git a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul.py b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul.py
index e3e381158e70f8..666744c4ad6006 100644
--- a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul.py
+++ b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul.py
@@ -19,8 +19,10 @@
 from absl import logging
 import immutabledict
 from jax.experimental.pallas import tpu as pltpu
+import jax.numpy as jnp
 
 from xla.benchmarks.core import benchmark
+from xla.benchmarks.core import platform_info
 from xla.benchmarks.pallas_microbenchmarks import subchannel_matmul_lib
 
 immutabledict = immutabledict.immutabledict
@@ -111,7 +113,12 @@ def main(_):
   lhs_dtype = benchmark.str_to_dtype(_FMT.value[0])
   rhs_dtype = benchmark.str_to_dtype(_FMT.value[1])
   out_dtype = benchmark.str_to_dtype(_FMT.value[2])
-  acc_dtype = benchmark.str_to_dtype(_ACC_DTYPE.value)
+  pinfo = platform_info.get_platform_info()
+  if jnp.issubdtype(lhs_dtype, jnp.integer) and pinfo.generation < 8:
+    inner_acc_dtype = jnp.int32
+  else:
+    inner_acc_dtype = jnp.float32
+  outer_acc_dtype = benchmark.str_to_dtype(_ACC_DTYPE.value)
 
   lhs_quantized_dtype = benchmark.str_to_dtype(_LHS_QUANTIZED_DTYPE.value)
   rhs_quantized_dtype = benchmark.str_to_dtype(_RHS_QUANTIZED_DTYPE.value)
@@ -149,7 +156,7 @@ def main(_):
         lhs_dtype=lhs_dtype,
         rhs_dtype=rhs_dtype,
         out_dtype=out_dtype,
-        acc_dtype=acc_dtype,
+        outer_acc_dtype=outer_acc_dtype,
         lhs_quantized_dtype=lhs_quantized_dtype,
         rhs_quantized_dtype=rhs_quantized_dtype,
         pre_quantize_lhs=pre_quantize_lhs,
@@ -185,7 +192,8 @@ def main(_):
       lhs_dtype=lhs_dtype,
       rhs_dtype=rhs_dtype,
       out_dtype=out_dtype,
-      acc_dtype=acc_dtype,
+      inner_acc_dtype=inner_acc_dtype,
+      outer_acc_dtype=outer_acc_dtype,
       lhs_quantized_dtype=lhs_quantized_dtype,
       rhs_quantized_dtype=rhs_quantized_dtype,
       pre_quantize_lhs=pre_quantize_lhs,
diff --git a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py
index 8a8f29fbbe6cdb..9bd532696b8e69 100644
--- a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py
+++ b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py
@@ -85,7 +85,7 @@ def select_window(
     lhs_dtype: jnp.dtype,
     rhs_dtype: jnp.dtype,
     out_dtype: jnp.dtype,
-    acc_dtype: jnp.dtype,
+    outer_acc_dtype: jnp.dtype,
     lhs_quantized_dtype: jnp.dtype,
     rhs_quantized_dtype: jnp.dtype,
     pre_quantize_lhs: bool = False,
@@ -116,7 +116,7 @@ def select_window(
       cost_lhs_dtype,
       cost_rhs_dtype,
       out_dtype,
-      acc_dtype,
+      outer_acc_dtype,
       chip_version=chip_version,
   ).select_window(
       vmem_limit_bytes,
@@ -143,7 +143,8 @@ class SubchannelMatmulConfig(benchmark.BenchmarkConfig):
     lhs_dtype: The unquantized dtype of the first operand.
     rhs_dtype: The unquantized dtype of the second operand.
     out_dtype: The dtype of the output.
-    acc_dtype: The dtype of the accumulator.
+    inner_acc_dtype: The dtype of the accumulator for the subchannel blocks.
+    outer_acc_dtype: The dtype of the accumulator for the output tiles.
     lhs_quantized_dtype: Quantized dtype of LHS.
     rhs_quantized_dtype: Quantized dtype of RHS.
     pre_quantize_lhs: Whether quantization of LHS happens outside kernel.
@@ -162,7 +163,8 @@ class SubchannelMatmulConfig(benchmark.BenchmarkConfig):
   lhs_dtype: jnp.dtype
   rhs_dtype: jnp.dtype
   out_dtype: jnp.dtype
-  acc_dtype: jnp.dtype
+  inner_acc_dtype: jnp.dtype
+  outer_acc_dtype: jnp.dtype
   lhs_quantized_dtype: jnp.dtype
   rhs_quantized_dtype: jnp.dtype
   pre_quantize_lhs: bool = False
@@ -180,7 +182,8 @@ def subchannel_matmul_kernel(
   block_m, block_k, block_n = cfg.block_m, cfg.block_k, cfg.block_n
   subchannel_size = cfg.subchannel_size
   lhs_mem, rhs_mem, out_mem = cfg.lhs_mem, cfg.rhs_mem, cfg.out_mem
-  acc_dtype = cfg.acc_dtype
+  inner_acc_dtype = cfg.inner_acc_dtype
+  outer_acc_dtype = cfg.outer_acc_dtype
 
   grid_m = math.ceil(m / block_m)
   grid_n = math.ceil(n / block_n)
@@ -219,14 +222,14 @@ def init():
         y_q_slice = rhs_ref[k_start:k_end, n_start:n_end]
 
         result = jnp.dot(
-            x_q_block, y_q_slice, preferred_element_type=jnp.float32
-        ).astype(acc_dtype)
+            x_q_block, y_q_slice, preferred_element_type=inner_acc_dtype
+        ).astype(outer_acc_dtype)
 
         y_safe_scales = rhs_scales_ref[sub_idx : sub_idx + 1, n_start:n_end]
         acc_ref[:, n_start:n_end] += (
             result
-            * x_safe_scales.astype(acc_dtype)
-            * y_safe_scales.astype(acc_dtype)
+            * x_safe_scales.astype(outer_acc_dtype)
+            * y_safe_scales.astype(outer_acc_dtype)
         )
 
     @pl.when(pl.program_id(2) == grid_k - 1)
@@ -291,7 +294,7 @@ def store():
           in_specs=in_specs,
           out_specs=out_specs,
           grid=(grid_m, grid_n, grid_k),
-          scratch_shapes=[pltpu.VMEM((block_m, block_n), acc_dtype)],
+          scratch_shapes=[pltpu.VMEM((block_m, block_n), outer_acc_dtype)],
       ),
       compiler_params=pltpu.CompilerParams(
           dimension_semantics=("parallel", "parallel", "arbitrary"),

From 1656f20935d9f090a071c9cf00c8a4a9530cadaf Mon Sep 17 00:00:00 2001
From: Matthias Guenther <mrguenther@google.com>
Date: Fri, 18 Sep 2026 14:40:26 -0700
Subject: [PATCH 07/13] Replace `/` with `_` in generated `lit_test` target
 names

This matches the behavior of some related tooling with which this needs to be compatible.

PiperOrigin-RevId: 984075260
---
 .../xla/third_party/stablehlo/temporary.patch | 1862 +++++++++++++++--
 1 file changed, 1724 insertions(+), 138 deletions(-)

diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch
index b6c9e2d8cb1f8a..ec90dce6dc5dc6 100644
--- a/third_party/xla/third_party/stablehlo/temporary.patch
+++ b/third_party/xla/third_party/stablehlo/temporary.patch
@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 #
-# This file is automatically generated by generate_patch tool.
+# This file is automatically generated by the `generate_patch` tool.
 # Do not edit directly.
 diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel
 --- stablehlo/BUILD.bazel
@@ -38,14 +38,14 @@ diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel
 +    ],
      licenses = ["notice"],
  )
-+
+ 
 +# copybara:uncomment_begin
 +# license(
 +#     name = "license",
 +#     package_name = "stablehlo",
 +# )
 +# copybara:uncomment_end
- 
++
  exports_files([
      "LICENSE",
      # Export to be used by MLIR-free python binding impls
@@ -323,78 +323,34 @@ diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel
      deps = [
          ":reference_errors",
          ":reference_types",
-@@ -588,38 +626,41 @@
+@@ -588,7 +626,8 @@
      hdrs = [
          "stablehlo/reference/Errors.h",
      ],
+-    strip_include_prefix = ".",
 +    compatible_with = get_compatible_with_portable(),
 +    includes = ["."],
-+    deps = [
-+        "@llvm-project//llvm:Support",
-+    ],
-+)
-+
-+cc_library(
-+    name = "reference_index",
-+    srcs = [
-+        "stablehlo/reference/Index.cpp",
-+    ],
-+    hdrs = [
-+        "stablehlo/reference/Index.h",
-+    ],
+     deps = [
+         "@llvm-project//llvm:Support",
+     ],
+@@ -602,7 +641,8 @@
+     hdrs = [
+         "stablehlo/reference/Index.h",
+     ],
+-    strip_include_prefix = ".",
 +    compatible_with = get_compatible_with_portable(),
 +    includes = ["."],
-+    deps = [
-+        "@llvm-project//llvm:Support",
-+        "@llvm-project//mlir:IR",
-+        "@llvm-project//mlir:Support",
-+    ],
-+)
-+
-+cc_library(
-+    name = "reference_numpy",
-+    srcs = [
-+        "stablehlo/reference/NumPy.cpp",
-+    ],
-+    hdrs = [
-+        "stablehlo/reference/NumPy.h",
-+    ],
+     deps = [
+         "@llvm-project//llvm:Support",
+         "@llvm-project//mlir:IR",
+@@ -618,6 +658,7 @@
+     hdrs = [
+         "stablehlo/reference/NumPy.h",
+     ],
 +    compatible_with = get_compatible_with_portable(),
      strip_include_prefix = ".",
      deps = [
--        "@llvm-project//llvm:Support",
--    ],
--)
--
--cc_library(
--    name = "reference_index",
--    srcs = [
--        "stablehlo/reference/Index.cpp",
--    ],
--    hdrs = [
--        "stablehlo/reference/Index.h",
--    ],
--    strip_include_prefix = ".",
--    deps = [
--        "@llvm-project//llvm:Support",
--        "@llvm-project//mlir:IR",
--        "@llvm-project//mlir:Support",
--    ],
--)
--
--cc_library(
--    name = "reference_numpy",
--    srcs = [
--        "stablehlo/reference/NumPy.cpp",
--    ],
--    hdrs = [
--        "stablehlo/reference/NumPy.h",
--    ],
--    strip_include_prefix = ".",
--    deps = [
          ":reference_tensor",
-         "@llvm-project//llvm:Support",
-         "@llvm-project//mlir:IR",
 @@ -635,7 +676,8 @@
      hdrs = [
          "stablehlo/reference/Ops.h",
@@ -942,7 +898,7 @@ diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel
      tbl_outs = {
          "stablehlo/dialect/VhloTypeDefs.cpp.inc": ["-gen-typedef-defs"],
          "stablehlo/dialect/VhloTypeDefs.h.inc": ["-gen-typedef-decls"],
-@@ -1826,227 +1896,7 @@
+@@ -1826,227 +1896,21 @@
      ],
  )
  
@@ -1152,8 +1108,11 @@ diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel
 -        "@llvm-project//mlir:InferTypeOpInterface",
 -        "@llvm-project//mlir:Support",
 -    ],
--)
--
++test_suite(
++    name = "all_tests",
++    tests = [":stablehlo_ci_tests"],
+ )
+ 
 -cc_test(
 -    name = "stablehlo_builder_test",
 -    srcs = ["stablehlo/integrations/cpp/builder/StablehloBuilderTest.cpp"],
@@ -1170,10 +1129,19 @@ diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel
 -        "@llvm-project//third-party/unittest:gtest",
 -    ],
 -)
-+test_suite(
-+    name = "all_tests",
-+    tests = [":stablehlo_ci_tests"],
-+)
++# copybara:uncomment_begin
++# alias(
++#     name = "check_ops",
++#     actual = "//stablehlo/tests:check_ops",
++#     deprecation = "Use //stablehlo/tests:check_ops instead.",
++# )
++#
++# alias(
++#     name = "test_utils",
++#     actual = "//stablehlo/tests:test_utils",
++#     deprecation = "Use //stablehlo/tests:test_utils instead.",
++# )
++# copybara:uncomment_end
 diff --ruN a/stablehlo/build_tools/BUILD.bazel b/stablehlo/build_tools/BUILD.bazel
 --- stablehlo/build_tools/BUILD.bazel
 +++ stablehlo/build_tools/BUILD.bazel
@@ -1219,7 +1187,7 @@ diff --ruN a/stablehlo/build_tools/BUILD.bazel b/stablehlo/build_tools/BUILD.baz
 diff --ruN a/stablehlo/build_tools/lit_test_suite.bzl b/stablehlo/build_tools/lit_test_suite.bzl
 --- stablehlo/build_tools/lit_test_suite.bzl
 +++ stablehlo/build_tools/lit_test_suite.bzl
-@@ -0,0 +1,69 @@
+@@ -0,0 +1,71 @@
 +# Copyright 2026 The StableHLO Authors. All Rights Reserved.
 +#
 +# Licensed under the Apache License, Version 2.0 (the "License");
@@ -1266,20 +1234,22 @@ diff --ruN a/stablehlo/build_tools/lit_test_suite.bzl b/stablehlo/build_tools/li
 +    tests_in_suite = []
 +
 +    for src in srcs:
-+        test_name = "%s.test" % src
++        test_name = ("%s.test" % src).replace("/", "_")  # copybara:comment
++
++        # copybara:uncomment test_name = "%s.test" % src
 +        lit_test(
-+            name = test_name,
-+            # name = src,  # copybara:comment
-+            srcs = [src],
-+            # src = src,  # copybara:comment
++            name = test_name,  # copybara:comment
++            # copybara:uncomment name = src,
++            srcs = [src],  # copybara:comment
++            # copybara:uncomment src = src,
 +            args = args,
 +            data = data + native.glob(
 +                ["%s.bc" % src],
 +                # Allow empty since not all `.mlir` tests have `.mlir.bc` files.
 +                allow_empty = True,
 +            ),
-+            deps = ["@rules_python//python/runfiles"],
-+            # driver = "//third_party/llvm/llvm-project/mlir:run_lit.sh",  # copybara:comment
++            deps = ["@rules_python//python/runfiles"],  # copybara:comment
++            # copybara:uncomment driver = "@llvm-project//mlir:run_lit.sh",
 +            tags = tags,
 +            size = size_per_test,
 +        )
@@ -1799,7 +1769,7 @@ diff --ruN a/stablehlo/examples/cpp/README.md b/stablehlo/examples/cpp/README.md
 diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel b/stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel
 --- stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel
 +++ stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel
-@@ -11,13 +11,26 @@
+@@ -11,14 +11,27 @@
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
@@ -1817,7 +1787,7 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel b/stablehl
 +    ],
      licenses = ["notice"],
  )
-+
+ 
 +# copybara:uncomment_begin
 +# filegroup(
 +#     name = "all_files",
@@ -1825,9 +1795,10 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel b/stablehl
 +#     visibility = ["//:__subpackages__"],
 +# )
 +# copybara:uncomment_end
- 
++
  # Equivalent of configure_lit_site_cfg from CMakeLists.txt.
  expand_template(
+     name = "lit_site_cfg_py_gen",
 @@ -39,25 +52,15 @@
      template = "lit.site.cfg.py.in",
  )
@@ -1909,10 +1880,69 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/reduce.mlir b/stablehl
  // CHECK-SAME:       strides = dense<2> : vector<2xi64>}
  // CHECK-SAME:      ins(%[[ARG0]], %[[WINDOW]] : tensor<1x17x17x64xf32>, tensor<3x3xf32>)
  // CHECK-SAME:      outs(%[[FILL]] : tensor<1x8x8x64xf32>) -> tensor<1x8x8x64xf32>
+diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h b/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h
+--- stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h
++++ stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h
+@@ -60,26 +60,26 @@
+     unsigned nParallelLoops);
+ 
+ /// Generates an init sparse tensor.
+-Value getEmptySparseTensor(OpBuilder &b, Location loc, ShapedType type,
++Value getEmptySparseTensor(OpBuilder& b, Location loc, ShapedType type,
+                            ArrayRef<Value> dynSizes);
+ 
+ /// Generates a tensor.empty op.
+-Value getEmptyTensor(OpBuilder &b, Location loc, ShapedType type,
++Value getEmptyTensor(OpBuilder& b, Location loc, ShapedType type,
+                      ArrayRef<Value> dynSizes);
+ 
+ /// Generates an empty tensor for the result of the operation, which could be a
+ /// dense tensor or a sparse tensor.
+-Value getEmptyTensorFor(OpBuilder &b, Location loc, ShapedType resultType,
+-                        Operation *op, ValueRange operands);
++Value getEmptyTensorFor(OpBuilder& b, Location loc, ShapedType resultType,
++                        Operation* op, ValueRange operands);
+ 
+ /// Ensures a tensor has the same shape (not including the element type) as
+ /// another.
+-Value coerceTensorShape(OpBuilder &builder, Location loc,
++Value coerceTensorShape(OpBuilder& builder, Location loc,
+                         TypedValue<ShapedType> value, ShapedType targetType);
+ 
+ /// Fills |tensor| with a zero constant of the matching type. Returns the new
+ /// value.
+-Value fillTensorWithZeros(OpBuilder &builder, Location loc, Value tensor);
++Value fillTensorWithZeros(OpBuilder& builder, Location loc, Value tensor);
+ 
+ /// Sparsifies a (block of) operation(s) that cannot be handled directly
+ /// by the sparse compiler but has well-known semi-ring semantics.
+@@ -94,17 +94,17 @@
+ ///     }
+ ///     absent={}
+ ///   linalg.yield %result
+-Value preSparsify(Operation *op, llvm::SmallVector<Value, 2> &values, Type rtp,
+-                  OpBuilder *b);
++Value preSparsify(Operation* op, llvm::SmallVector<Value, 2>& values, Type rtp,
++                  OpBuilder* b);
+ 
+ /// Finalizes sparse semi-ring construction.
+-Value postSparsify(Operation *op, Value semiring, Value result, OpBuilder *b);
++Value postSparsify(Operation* op, Value semiring, Value result, OpBuilder* b);
+ 
+ /// Returns true if all operands are tensors with rank 0.
+-bool allOperandsAreScalarTensors(Operation *op);
++bool allOperandsAreScalarTensors(Operation* op);
+ 
+ /// Returns true if parent op is linalg.
+-bool isInBodyOfLinalgOps(Operation *op);
++bool isInBodyOfLinalgOps(Operation* op);
+ 
+ /// Extracts integer values from the attribute |elements|.
+ SmallVector<int64_t> extract1DVector(DenseIntElementsAttr elements);
 diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel b/stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel
 --- stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel
 +++ stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel
-@@ -11,13 +11,26 @@
+@@ -11,14 +11,27 @@
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
@@ -1930,7 +1960,7 @@ diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel b/stablehlo/
 +    ],
      licenses = ["notice"],
  )
-+
+ 
 +# copybara:uncomment_begin
 +# filegroup(
 +#     name = "all_files",
@@ -1938,9 +1968,10 @@ diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel b/stablehlo/
 +#     visibility = ["//:__subpackages__"],
 +# )
 +# copybara:uncomment_end
- 
++
  # Equivalent of configure_lit_site_cfg from CMakeLists.txt.
  expand_template(
+     name = "lit_site_cfg_py_gen",
 @@ -39,26 +52,16 @@
      template = "lit.site.cfg.py.in",
  )
@@ -1978,6 +2009,147 @@ diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel b/stablehlo/
 +    size_per_test = "small",
      tags = ["stablehlo_tosa_tests"],
  )
+diff --ruN a/stablehlo/stablehlo/dialect/ChloBytecode.cpp b/stablehlo/stablehlo/dialect/ChloBytecode.cpp
+--- stablehlo/stablehlo/dialect/ChloBytecode.cpp
++++ stablehlo/stablehlo/dialect/ChloBytecode.cpp
+@@ -111,37 +111,37 @@
+ /// This class implements the bytecode interface for the  CHLO dialect.
+ class ChloBytecodeInterface : public BytecodeDialectInterface {
+  public:
+-  ChloBytecodeInterface(Dialect *dialect) : BytecodeDialectInterface(dialect) {}
++  ChloBytecodeInterface(Dialect* dialect) : BytecodeDialectInterface(dialect) {}
+ 
+   //===--------------------------------------------------------------------===//
+   // Attributes
+ 
+   // These methods are invoked by superclass when an attr from  CHLO dialect
+   // is encountered.
+-  Attribute readAttribute(DialectBytecodeReader &reader) const override;
++  Attribute readAttribute(DialectBytecodeReader& reader) const override;
+   LogicalResult writeAttribute(Attribute attr,
+-                               DialectBytecodeWriter &writer) const override;
++                               DialectBytecodeWriter& writer) const override;
+ 
+   // TO ADD ATTRIBUTE: Include a read method for each attribute in CHLO
+   // Ex: SomeAttr readSomeAttr(DialectBytecodeReader &reader) const;
+   ComparisonDirectionAttr readComparisonDirectionAttr(
+-      DialectBytecodeReader &reader) const;
++      DialectBytecodeReader& reader) const;
+   ComparisonTypeAttr readComparisonTypeAttr(
+-      DialectBytecodeReader &reader) const;
++      DialectBytecodeReader& reader) const;
+ 
+   // TO ADD ATTRIBUTE: Include a write method for each attribute in CHLO
+   // Ex: void write(SomeAttr attr, DialectBytecodeWriter &writer) const;
+-  void write(ComparisonDirectionAttr attr, DialectBytecodeWriter &writer) const;
+-  void write(ComparisonTypeAttr attr, DialectBytecodeWriter &writer) const;
++  void write(ComparisonDirectionAttr attr, DialectBytecodeWriter& writer) const;
++  void write(ComparisonTypeAttr attr, DialectBytecodeWriter& writer) const;
+ 
+   //===--------------------------------------------------------------------===//
+   // Types
+ 
+   // These methods are invoked by superclass when a type from  CHLO dialect
+   // is encountered.
+-  Type readType(DialectBytecodeReader &reader) const override;
++  Type readType(DialectBytecodeReader& reader) const override;
+   LogicalResult writeType(Type type,
+-                          DialectBytecodeWriter &writer) const override;
++                          DialectBytecodeWriter& writer) const override;
+ 
+   // TO ADD TYPE: Include a read method for each type in CHLO
+   // Ex: SomeType readSomeType(DialectBytecodeReader &reader) const;
+@@ -158,7 +158,7 @@
+ 
+ // TO ADD ATTRIBUTE: Update the switch to include a branch for the attr.
+ Attribute ChloBytecodeInterface::readAttribute(
+-    DialectBytecodeReader &reader) const {
++    DialectBytecodeReader& reader) const {
+   uint64_t code;
+   if (failed(reader.readVarInt(code))) return Attribute();
+   switch (code) {
+@@ -173,7 +173,7 @@
+ }
+ 
+ ComparisonDirectionAttr ChloBytecodeInterface::readComparisonDirectionAttr(
+-    DialectBytecodeReader &reader) const {
++    DialectBytecodeReader& reader) const {
+   LOG_READ_CALL;
+   return hlo::bytecode::readEnumAttribute<ComparisonDirectionAttr>(
+       reader, getContext(),
+@@ -181,7 +181,7 @@
+ }
+ 
+ ComparisonTypeAttr ChloBytecodeInterface::readComparisonTypeAttr(
+-    DialectBytecodeReader &reader) const {
++    DialectBytecodeReader& reader) const {
+   LOG_READ_CALL;
+   return hlo::bytecode::readEnumAttribute<ComparisonTypeAttr>(
+       reader, getContext(),
+@@ -195,7 +195,7 @@
+ // If this method returns failure, the string serialization is used in the
+ // bytecode.
+ LogicalResult ChloBytecodeInterface::writeAttribute(
+-    Attribute attr, DialectBytecodeWriter &writer) const {
++    Attribute attr, DialectBytecodeWriter& writer) const {
+   return TypeSwitch<Attribute, LogicalResult>(attr)
+       .Case<ComparisonDirectionAttr, ComparisonTypeAttr>([&](auto attr) {
+         LOG_WRITE_CALL;
+@@ -209,13 +209,13 @@
+ }
+ 
+ void ChloBytecodeInterface::write(ComparisonDirectionAttr attr,
+-                                  DialectBytecodeWriter &writer) const {
++                                  DialectBytecodeWriter& writer) const {
+   writer.writeVarInt(chlo_encoding::kComparisonDirectionAttr);
+   hlo::bytecode::writeEnumAttribute<ComparisonDirection>(attr, writer);
+ }
+ 
+ void ChloBytecodeInterface::write(ComparisonTypeAttr attr,
+-                                  DialectBytecodeWriter &writer) const {
++                                  DialectBytecodeWriter& writer) const {
+   writer.writeVarInt(chlo_encoding::kComparisonTypeAttr);
+   hlo::bytecode::writeEnumAttribute<ComparisonType>(attr, writer);
+ }
+@@ -224,7 +224,7 @@
+ // Types: Reader
+ 
+ // TO ADD TYPE: Update the case selection to include the new type.
+-Type ChloBytecodeInterface::readType(DialectBytecodeReader &reader) const {
++Type ChloBytecodeInterface::readType(DialectBytecodeReader& reader) const {
+   uint64_t code;
+   if (failed(reader.readVarInt(code))) return Type();
+ 
+@@ -240,7 +240,7 @@
+ 
+ // TO ADD TYPE: Update the case selection to include the new type.
+ LogicalResult ChloBytecodeInterface::writeType(
+-    Type type, DialectBytecodeWriter &writer) const {
++    Type type, DialectBytecodeWriter& writer) const {
+   return TypeSwitch<Type, LogicalResult>(type).Default([&](Type) {
+     LOG_NOT_IMPLEMENTED;
+     return failure();
+@@ -249,7 +249,7 @@
+ 
+ }  // namespace
+ 
+-void addBytecodeInterface(ChloDialect *dialect) {
++void addBytecodeInterface(ChloDialect* dialect) {
+   dialect->addInterfaces<ChloBytecodeInterface>();
+ }
+ }  // namespace chlo
+diff --ruN a/stablehlo/stablehlo/dialect/ChloBytecode.h b/stablehlo/stablehlo/dialect/ChloBytecode.h
+--- stablehlo/stablehlo/dialect/ChloBytecode.h
++++ stablehlo/stablehlo/dialect/ChloBytecode.h
+@@ -21,7 +21,7 @@
+ 
+ // Add the interface necessary for encoding and decoding the CHLO dialect
+ // components in bytecode.
+-void addBytecodeInterface(ChloDialect *dialect);
++void addBytecodeInterface(ChloDialect* dialect);
+ }  // namespace chlo
+ }  // namespace mlir
+ 
 diff --ruN a/stablehlo/stablehlo/dialect/ChloEnums.td b/stablehlo/stablehlo/dialect/ChloEnums.td
 --- stablehlo/stablehlo/dialect/ChloEnums.td
 +++ stablehlo/stablehlo/dialect/ChloEnums.td
@@ -2031,11 +2203,10 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloOps.cpp b/stablehlo/stablehlo/diale
    for (auto [i, type] : llvm::enumerate(adaptor.getInputs().getTypes())) {
      auto inputType = dyn_cast<RankedTensorType>(type);
      if (!inputType) {
-@@ -839,7 +843,26 @@
+@@ -839,8 +843,27 @@
      if (dimSize == ShapedType::kDynamic) {
        dimSize = inputType.getDimSize(dim);
      }
--  }
 +    if (inputType.getEncoding()) {
 +      prototypeEncoding = inputType.getEncoding();
 +    }
@@ -2050,15 +2221,16 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloOps.cpp b/stablehlo/stablehlo/diale
 +        }
 +      }
 +    }
-+  }
-+
+   }
+ 
 +  // If the scan dimension size is dynamic, propagate any bound inferred from
 +  // the inputs; otherwise static dimensions do not need a bound.
 +  const int64_t insertBound =
 +      (dimSize == ShapedType::kDynamic) ? dimBound : ShapedType::kDynamic;
- 
++
    for (auto [i, type] : llvm::enumerate(terminator->getOperands().getTypes())) {
      auto resultType = dyn_cast<RankedTensorType>(type);
+     if (!resultType) {
 @@ -849,10 +872,34 @@
      }
      SmallVector<int64_t> shape(resultType.getShape().begin(),
@@ -2069,8 +2241,6 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloOps.cpp b/stablehlo/stablehlo/diale
 +    // carries whose shapes and encodings are preserved as-is.
      if (i < numOutputs) {
        shape.insert(std::next(shape.begin(), dim), dimSize);
--    }
--    inferredReturnShapes.emplace_back(shape, resultType.getElementType());
 +      SmallVector<int64_t> resultBounds =
 +          llvm::to_vector(hlo::encodingToBounds(encoding));
 +      // Re-insert bounds along the scan dimension if the result has bounds or
@@ -2090,7 +2260,8 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloOps.cpp b/stablehlo/stablehlo/diale
 +          encoding = hlo::boundsToEncoding(encoding, resultBounds);
 +        }
 +      }
-+    }
+     }
+-    inferredReturnShapes.emplace_back(shape, resultType.getElementType());
 +    inferredReturnShapes.emplace_back(shape, resultType.getElementType(),
 +                                      encoding);
    }
@@ -2115,6 +2286,30 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloOps.cpp b/stablehlo/stablehlo/diale
      }
      if (!hlo::isCompatibleForHloTypeInference(
              argType, bodyBlock.getArgument(i).getType())) {
+diff --ruN a/stablehlo/stablehlo/dialect/Register.cpp b/stablehlo/stablehlo/dialect/Register.cpp
+--- stablehlo/stablehlo/dialect/Register.cpp
++++ stablehlo/stablehlo/dialect/Register.cpp
+@@ -27,7 +27,7 @@
+ namespace mlir {
+ namespace stablehlo {
+ 
+-void registerAllDialects(mlir::DialectRegistry &registry) {
++void registerAllDialects(mlir::DialectRegistry& registry) {
+   // clang-format off
+   registry.insert<mlir::func::FuncDialect,
+                   mlir::quant::QuantDialect,
+diff --ruN a/stablehlo/stablehlo/dialect/Register.h b/stablehlo/stablehlo/dialect/Register.h
+--- stablehlo/stablehlo/dialect/Register.h
++++ stablehlo/stablehlo/dialect/Register.h
+@@ -23,7 +23,7 @@
+ namespace stablehlo {
+ 
+ // Add chlo and stablehlo dialects to the provided registry.
+-void registerAllDialects(DialectRegistry &registry);
++void registerAllDialects(DialectRegistry& registry);
+ 
+ }  // namespace stablehlo
+ }  // namespace mlir
 diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp
 --- stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp
 +++ stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp
@@ -2169,15 +2364,13 @@ diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stabl
 -        break;
 -      }
 -    }
--  }
--
--  SmallVector<int64_t> transposedSizes(meshAxisNames.size());
--  for (size_t i = 0; i < meshAxisNames.size(); ++i) {
 +  // Grouped axes in the specified order
 +  for (int64_t idx : groupedAxisIndices) {
 +    transposeAxes.push_back(idx);
-+  }
-+
+   }
+ 
+-  SmallVector<int64_t> transposedSizes(meshAxisNames.size());
+-  for (size_t i = 0; i < meshAxisNames.size(); ++i) {
 +  SmallVector<int64_t> transposedSizes(axisSizes.size());
 +  for (size_t i = 0; i < axisSizes.size(); ++i) {
      transposedSizes[i] = axisSizes[transposeAxes[i]];
@@ -2223,14 +2416,14 @@ diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stabl
          int64_t coord =
              (linearTransposeIdx / transposedStrides[k]) % transposedSizes[k];
          originalIndex += coord * originalStrides[transposeAxes[k]];
-@@ -102,9 +99,135 @@
+@@ -102,10 +99,136 @@
      }
      groups.push_back(std::move(group));
    }
 -
    return groups;
  }
-+
+ 
 +// Splits mesh axes based on sub-axis references and computes the corresponding
 +// indices for the communication axes.
 +FailureOr<ReindexedAxes> computeReindexedAxes(ArrayRef<MeshAxisAttr> axesInMesh,
@@ -2357,9 +2550,10 @@ diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stabl
 +}
 +
 +}  // namespace
- 
++
  FailureOr<SmallVector<SmallVector<int64_t>>> flattenReplicaGroupMeshAxes(
      Attribute meshAttr, ArrayAttr commAxes, Location loc) {
+   // Reference implementation for flattening a mesh-axes based replica group
 @@ -120,34 +243,13 @@
    if (!mesh)
      return emitOptionalError(loc, "expected stablehlo.mesh for mesh attribute");
@@ -2411,6 +2605,18 @@ diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stabl
  }
  
  }  // namespace stablehlo
+diff --ruN a/stablehlo/stablehlo/dialect/StablehloBytecode.h b/stablehlo/stablehlo/dialect/StablehloBytecode.h
+--- stablehlo/stablehlo/dialect/StablehloBytecode.h
++++ stablehlo/stablehlo/dialect/StablehloBytecode.h
+@@ -21,7 +21,7 @@
+ 
+ // Add the interface necessary for encoding and decoding StableHLO dialect
+ // components in bytecode.
+-void addBytecodeInterface(StablehloDialect *dialect);
++void addBytecodeInterface(StablehloDialect* dialect);
+ }  // namespace stablehlo
+ }  // namespace mlir
+ 
 diff --ruN a/stablehlo/stablehlo/dialect/StablehloEnums.td b/stablehlo/stablehlo/dialect/StablehloEnums.td
 --- stablehlo/stablehlo/dialect/StablehloEnums.td
 +++ stablehlo/stablehlo/dialect/StablehloEnums.td
@@ -2491,6 +2697,18 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloEnums.td b/stablehlo/stablehlo
  
  #endif // STABLEHLO_DIALECT_STABLEHLO_ENUMS
 +
+diff --ruN a/stablehlo/stablehlo/dialect/VhloBytecode.h b/stablehlo/stablehlo/dialect/VhloBytecode.h
+--- stablehlo/stablehlo/dialect/VhloBytecode.h
++++ stablehlo/stablehlo/dialect/VhloBytecode.h
+@@ -21,7 +21,7 @@
+ 
+ // Add the interface necessary for encoding and decoding VHLO dialect
+ // components in bytecode.
+-void addBytecodeInterface(VhloDialect *dialect);
++void addBytecodeInterface(VhloDialect* dialect);
+ }  // namespace vhlo
+ }  // namespace mlir
+ 
 diff --ruN a/stablehlo/stablehlo/dialect/VhloEnums.td b/stablehlo/stablehlo/dialect/VhloEnums.td
 --- stablehlo/stablehlo/dialect/VhloEnums.td
 +++ stablehlo/stablehlo/dialect/VhloEnums.td
@@ -2502,6 +2720,169 @@ diff --ruN a/stablehlo/stablehlo/dialect/VhloEnums.td b/stablehlo/stablehlo/dial
    let extraClassDeclaration = [{
      mlir::vhlo::Version getMinVersion() {
        return mlir::vhlo::Version(}] # !subst(".", ", ", minVersion) # [{);
+diff --ruN a/stablehlo/stablehlo/dialect/VhloOps.h b/stablehlo/stablehlo/dialect/VhloOps.h
+--- stablehlo/stablehlo/dialect/VhloOps.h
++++ stablehlo/stablehlo/dialect/VhloOps.h
+@@ -31,20 +31,20 @@
+ 
+ class VhloDialect : public Dialect {
+  public:
+-  explicit VhloDialect(MLIRContext *context);
++  explicit VhloDialect(MLIRContext* context);
+   static StringRef getDialectNamespace() { return "vhlo"; }
+ 
+   // Parses a type registered to this dialect.
+-  Type parseType(DialectAsmParser &parser) const override;
++  Type parseType(DialectAsmParser& parser) const override;
+ 
+   // Prints a type registered to this dialect.
+-  void printType(Type type, DialectAsmPrinter &os) const override;
++  void printType(Type type, DialectAsmPrinter& os) const override;
+ 
+   // Parses an attribute registered to this dialect.
+-  Attribute parseAttribute(DialectAsmParser &parser, Type type) const override;
++  Attribute parseAttribute(DialectAsmParser& parser, Type type) const override;
+ 
+   // Prints an attribute registered to this dialect.
+-  void printAttribute(Attribute attr, DialectAsmPrinter &os) const override;
++  void printAttribute(Attribute attr, DialectAsmPrinter& os) const override;
+ 
+  private:
+   // Adds VHLO types to this dialect.
+diff --ruN a/stablehlo/stablehlo/integrations/c/ChloAttributes.cpp b/stablehlo/stablehlo/integrations/c/ChloAttributes.cpp
+--- stablehlo/stablehlo/integrations/c/ChloAttributes.cpp
++++ stablehlo/stablehlo/integrations/c/ChloAttributes.cpp
+@@ -73,12 +73,12 @@
+ 
+ MlirAttribute chloRaggedDotDimensionNumbersGet(
+     MlirContext ctx, intptr_t nLhsBatchingDimensions,
+-    const int64_t *lhsBatchingDimensions, intptr_t nRhsBatchingDimensions,
+-    const int64_t *rhsBatchingDimensions, intptr_t nLhsContractingDimensions,
+-    const int64_t *lhsContractingDimensions, intptr_t nRhsContractingDimensions,
+-    const int64_t *rhsContractingDimensions, intptr_t nLhsRaggedDimensions,
+-    const int64_t *lhsRaggedDimensions, intptr_t nRhsGroupDimensions,
+-    const int64_t *rhsGroupDimensions) {
++    const int64_t* lhsBatchingDimensions, intptr_t nRhsBatchingDimensions,
++    const int64_t* rhsBatchingDimensions, intptr_t nLhsContractingDimensions,
++    const int64_t* lhsContractingDimensions, intptr_t nRhsContractingDimensions,
++    const int64_t* rhsContractingDimensions, intptr_t nLhsRaggedDimensions,
++    const int64_t* lhsRaggedDimensions, intptr_t nRhsGroupDimensions,
++    const int64_t* rhsGroupDimensions) {
+   return wrap(mlir::chlo::RaggedDotDimensionNumbersAttr::get(
+       unwrap(ctx),
+       llvm::ArrayRef(lhsBatchingDimensions, nLhsBatchingDimensions),
+diff --ruN a/stablehlo/stablehlo/integrations/c/ChloAttributes.h b/stablehlo/stablehlo/integrations/c/ChloAttributes.h
+--- stablehlo/stablehlo/integrations/c/ChloAttributes.h
++++ stablehlo/stablehlo/integrations/c/ChloAttributes.h
+@@ -53,16 +53,16 @@
+ 
+ MLIR_CAPI_EXPORTED MlirAttribute chloRaggedDotDimensionNumbersGet(
+     MlirContext ctx,                                                        //
+-    intptr_t nLhsBatchingDimensions, const int64_t *lhsBatchingDimensions,  //
+-    intptr_t nRhsBatchingDimensions, const int64_t *rhsBatchingDimensions,  //
++    intptr_t nLhsBatchingDimensions, const int64_t* lhsBatchingDimensions,  //
++    intptr_t nRhsBatchingDimensions, const int64_t* rhsBatchingDimensions,  //
+     intptr_t nLhsContractingDimensions,                                     //
+-    const int64_t *lhsContractingDimensions,                                //
++    const int64_t* lhsContractingDimensions,                                //
+     intptr_t nRhsContractingDimensions,                                     //
+-    const int64_t *rhsContractingDimensions,                                //
++    const int64_t* rhsContractingDimensions,                                //
+     intptr_t nLhsRaggedDimensions,                                          //
+-    const int64_t *lhsRaggedDimensions,                                     //
++    const int64_t* lhsRaggedDimensions,                                     //
+     intptr_t nRhsGroupDimensions,                                           //
+-    const int64_t *rhsGroupDimensions);
++    const int64_t* rhsGroupDimensions);
+ 
+ MLIR_CAPI_EXPORTED bool chloAttributeIsARaggedDotDimensionNumbers(
+     MlirAttribute attr);
+diff --ruN a/stablehlo/stablehlo/integrations/c/StablehloDialectApi.cpp b/stablehlo/stablehlo/integrations/c/StablehloDialectApi.cpp
+--- stablehlo/stablehlo/integrations/c/StablehloDialectApi.cpp
++++ stablehlo/stablehlo/integrations/c/StablehloDialectApi.cpp
+@@ -48,18 +48,18 @@
+ 
+ void stablehloVersionFromCompatibilityRequirement(
+     MlirStablehloCompatibilityRequirement requirement,
+-    MlirStringCallback callback, void *userData) {
++    MlirStringCallback callback, void* userData) {
+   mlir::detail::CallbackOstream stream(callback, userData);
+   stream << mlir::vhlo::Version::fromCompatibilityRequirement(
+       unwrapCompatibilityRequirement(requirement));
+ }
+ 
+-void stablehloGetCurrentVersion(MlirStringCallback callback, void *userData) {
++void stablehloGetCurrentVersion(MlirStringCallback callback, void* userData) {
+   mlir::detail::CallbackOstream stream(callback, userData);
+   stream << mlir::stablehlo::getCurrentVersion();
+ }
+ 
+-void stablehloGetMinimumVersion(MlirStringCallback callback, void *userData) {
++void stablehloGetMinimumVersion(MlirStringCallback callback, void* userData) {
+   mlir::detail::CallbackOstream stream(callback, userData);
+   stream << mlir::stablehlo::getMinimumVersion();
+ }
+@@ -67,7 +67,7 @@
+ MlirLogicalResult stablehloGetSmallerVersion(MlirStringRef version1,
+                                              MlirStringRef version2,
+                                              MlirStringCallback callback,
+-                                             void *userData) {
++                                             void* userData) {
+   mlir::detail::CallbackOstream stream(callback, userData);
+   auto result =
+       mlir::stablehlo::getSmallerVersion(unwrap(version1), unwrap(version2));
+@@ -78,7 +78,7 @@
+ 
+ MlirLogicalResult stablehloSerializePortableArtifactFromModule(
+     MlirModule moduleStr, MlirStringRef targetVersion,
+-    MlirStringCallback callback, void *userData, bool allowOtherDialects) {
++    MlirStringCallback callback, void* userData, bool allowOtherDialects) {
+   mlir::detail::CallbackOstream stream(callback, userData);
+   if (failed(mlir::stablehlo::serializePortableArtifact(
+           unwrap(moduleStr), unwrap(targetVersion), stream,
+@@ -89,7 +89,7 @@
+ 
+ MlirLogicalResult stablehloSerializePortableArtifactFromStringRef(
+     MlirStringRef moduleStr, MlirStringRef targetVersion,
+-    MlirStringCallback callback, void *userData) {
++    MlirStringCallback callback, void* userData) {
+   mlir::detail::CallbackOstream stream(callback, userData);
+   if (failed(mlir::stablehlo::serializePortableArtifact(
+           unwrap(moduleStr), unwrap(targetVersion), stream)))
+@@ -98,7 +98,7 @@
+ }
+ 
+ MlirLogicalResult stablehloDeserializePortableArtifact(
+-    MlirStringRef artifactStr, MlirStringCallback callback, void *userData) {
++    MlirStringRef artifactStr, MlirStringCallback callback, void* userData) {
+   mlir::detail::CallbackOstream stream(callback, userData);
+   if (failed(mlir::stablehlo::deserializePortableArtifact(unwrap(artifactStr),
+                                                           stream)))
+diff --ruN a/stablehlo/stablehlo/integrations/c/StablehloUnifiedApi.cpp b/stablehlo/stablehlo/integrations/c/StablehloUnifiedApi.cpp
+--- stablehlo/stablehlo/integrations/c/StablehloUnifiedApi.cpp
++++ stablehlo/stablehlo/integrations/c/StablehloUnifiedApi.cpp
+@@ -32,9 +32,9 @@
+ #include "stablehlo/reference/Configuration.h"
+ 
+ MlirAttribute stablehloEvalModule(MlirModule module, int nArgs,
+-                                  MlirAttribute const *args,
+-                                  const char *const probeInstrumentationDir,
+-                                  int *errorCode) {
++                                  MlirAttribute const* args,
++                                  const char* const probeInstrumentationDir,
++                                  int* errorCode) {
+   std::vector<mlir::DenseElementsAttr> inputs;
+   inputs.reserve(nArgs);
+   for (int i = 0; i < nArgs; ++i) {
+@@ -49,7 +49,7 @@
+     return MlirAttribute{nullptr};
+   }
+   std::vector<MlirAttribute> resultsVec;
+-  for (const auto &result : results.value()) {
++  for (const auto& result : results.value()) {
+     resultsVec.push_back(wrap(result));
+   }
+   return mlirArrayAttrGet(mlirModuleGetContext(module), resultsVec.size(),
 diff --ruN a/stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel b/stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel
 --- stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel
 +++ stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel
@@ -2811,10 +3192,745 @@ diff --ruN a/stablehlo/stablehlo/integrations/cpp/builder/MlirBuilder.cpp b/stab
  
  #include "mlir/IR/Location.h"
  #include "mlir/IR/Value.h"
+diff --ruN a/stablehlo/stablehlo/integrations/python/StablehloApi.cpp b/stablehlo/stablehlo/integrations/python/StablehloApi.cpp
+--- stablehlo/stablehlo/integrations/python/StablehloApi.cpp
++++ stablehlo/stablehlo/integrations/python/StablehloApi.cpp
+@@ -42,15 +42,15 @@
+   StringWriterHelper() : ss_(s_) {}
+ 
+   static MlirStringCallback getMlirStringCallback() {
+-    return [](MlirStringRef string_ref, void *user_data) {
+-      auto *helper = static_cast<StringWriterHelper *>(user_data);
++    return [](MlirStringRef string_ref, void* user_data) {
++      auto* helper = static_cast<StringWriterHelper*>(user_data);
+       helper->ss_ << llvm::StringRef(string_ref.data, string_ref.length);
+     };
+   }
+ 
+-  void *getUserData() { return static_cast<void *>(this); }
++  void* getUserData() { return static_cast<void*>(this); }
+ 
+-  const std::string &toString() {
++  const std::string& toString() {
+     ss_.flush();
+     return s_;
+   }
+@@ -60,7 +60,7 @@
+   llvm::raw_string_ostream ss_;
+ };
+ 
+-static MlirStringRef toMlirStringRef(const std::string &s) {
++static MlirStringRef toMlirStringRef(const std::string& s) {
+   return mlirStringRefCreate(s.data(), s.size());
+ }
+ 
+@@ -68,11 +68,11 @@
+   return mlirStringRefCreate(s.data(), s.size());
+ }
+ 
+-static MlirStringRef toMlirStringRef(const nb::bytes &s) {
+-  return mlirStringRefCreate(static_cast<const char *>(s.data()), s.size());
++static MlirStringRef toMlirStringRef(const nb::bytes& s) {
++  return mlirStringRefCreate(static_cast<const char*>(s.data()), s.size());
+ }
+ 
+-void AddStablehloApi(nb::module_ &m) {
++void AddStablehloApi(nb::module_& m) {
+   // Portable API is a subset of StableHLO API
+   AddPortableApi(m);
+ 
+@@ -146,8 +146,8 @@
+   //
+   m.def(
+       "eval_module",
+-      [](MlirModule module, std::vector<MlirAttribute> &args,
+-         const std::string &probe_instrumentation_dir)
++      [](MlirModule module, std::vector<MlirAttribute>& args,
++         const std::string& probe_instrumentation_dir)
+           -> std::vector<MlirAttribute> {
+         for (auto arg : args) {
+           if (!mlirAttributeIsADenseElements(arg)) {
+@@ -174,7 +174,7 @@
+       nb::arg("probe_instrumentation_dir") = "");
+ }
+ 
+-void AddPortableApi(nb::module_ &m) {
++void AddPortableApi(nb::module_& m) {
+   //
+   // Utility APIs.
+   //
+@@ -182,8 +182,8 @@
+ 
+   m.def(
+       "get_smaller_version",
+-      [](const std::string &version1,
+-         const std::string &version2) -> std::string {
++      [](const std::string& version1,
++         const std::string& version2) -> std::string {
+         StringWriterHelper accumulator;
+         if (mlirLogicalResultIsFailure(stablehloGetSmallerVersion(
+                 toMlirStringRef(version1), toMlirStringRef(version2),
+@@ -263,7 +263,7 @@
+       nb::arg("artifact_str"));
+   m.def(
+       "deserialize_portable_artifact_str",
+-      [](const nb::bytes &artifact) -> nb::bytes {
++      [](const nb::bytes& artifact) -> nb::bytes {
+         StringWriterHelper accumulator;
+         if (mlirLogicalResultIsFailure(stablehloDeserializePortableArtifact(
+                 toMlirStringRef(artifact), accumulator.getMlirStringCallback(),
+diff --ruN a/stablehlo/stablehlo/reference/Api.cpp b/stablehlo/stablehlo/reference/Api.cpp
+--- stablehlo/stablehlo/reference/Api.cpp
++++ stablehlo/stablehlo/reference/Api.cpp
+@@ -83,7 +83,7 @@
+ class DefaultInterpreterFallback : public InterpreterFallback {
+  public:
+   DefaultInterpreterFallback(const InterpreterConfiguration& config)
+-      : config(config){};
++      : config(config) {};
+ 
+   virtual llvm::Error operator()(Operation& op, Scope& scope,
+                                  Process* process) final {
+diff --ruN a/stablehlo/stablehlo/reference/Api.h b/stablehlo/stablehlo/reference/Api.h
+--- stablehlo/stablehlo/reference/Api.h
++++ stablehlo/stablehlo/reference/Api.h
+@@ -36,17 +36,17 @@
+ /// builtin kernels are matched.
+ FailureOr<SmallVector<InterpreterValue>> evalModule(
+     ModuleOp module, ArrayRef<InterpreterValue> inputs,
+-    const InterpreterConfiguration &config);
++    const InterpreterConfiguration& config);
+ 
+ /// This wrapper is intended to be easily used by the StableHLO Python bindings.
+ // It wraps the InterpreterValue API.
+ FailureOr<SmallVector<DenseElementsAttr>> evalModule(
+     ModuleOp module, ArrayRef<DenseElementsAttr> inputs,
+-    const InterpreterConfiguration &config);
++    const InterpreterConfiguration& config);
+ 
+ /// Parses a StableHLO MLIR text program into a ModuleOp.
+-FailureOr<OwningOpRef<ModuleOp>> parseStablehloModule(const std::string &mlir,
+-                                                      MLIRContext &context);
++FailureOr<OwningOpRef<ModuleOp>> parseStablehloModule(const std::string& mlir,
++                                                      MLIRContext& context);
+ 
+ }  // namespace stablehlo
+ }  // namespace mlir
+diff --ruN a/stablehlo/stablehlo/reference/Axes.cpp b/stablehlo/stablehlo/reference/Axes.cpp
+--- stablehlo/stablehlo/reference/Axes.cpp
++++ stablehlo/stablehlo/reference/Axes.cpp
+@@ -21,7 +21,7 @@
+ namespace mlir {
+ namespace stablehlo {
+ 
+-raw_ostream &operator<<(raw_ostream &os, const Axes &x) {
++raw_ostream& operator<<(raw_ostream& os, const Axes& x) {
+   os << "[";
+   llvm::interleave(x, os, ", ");
+   os << "]";
+diff --ruN a/stablehlo/stablehlo/reference/Axes.h b/stablehlo/stablehlo/reference/Axes.h
+--- stablehlo/stablehlo/reference/Axes.h
++++ stablehlo/stablehlo/reference/Axes.h
+@@ -29,8 +29,8 @@
+ class Axes : public SmallVector<int64_t> {
+  public:
+   Axes() = default;
+-  Axes(const Axes &other) = default;
+-  Axes &operator=(const Axes &other) = default;
++  Axes(const Axes& other) = default;
++  Axes& operator=(const Axes& other) = default;
+ 
+   Axes(std::initializer_list<int64_t> list) : SmallVector(list) {}
+   explicit Axes(size_t size, int64_t element = 0)
+@@ -40,7 +40,7 @@
+       : SmallVector(attr.getValues<int64_t>()) {}
+ };
+ 
+-raw_ostream &operator<<(raw_ostream &os, const Axes &x);
++raw_ostream& operator<<(raw_ostream& os, const Axes& x);
+ 
+ }  // namespace stablehlo
+ }  // namespace mlir
+diff --ruN a/stablehlo/stablehlo/reference/Configuration.cpp b/stablehlo/stablehlo/reference/Configuration.cpp
+--- stablehlo/stablehlo/reference/Configuration.cpp
++++ stablehlo/stablehlo/reference/Configuration.cpp
+@@ -25,8 +25,8 @@
+ namespace mlir {
+ namespace stablehlo {
+ 
+-llvm::Error InterpreterFallback::operator()(Operation &op, Scope &scope,
+-                                            Process *process) {
++llvm::Error InterpreterFallback::operator()(Operation& op, Scope& scope,
++                                            Process* process) {
+   return stablehlo::invalidArgument("Unsupported op: %s",
+                                     debugString(op).c_str());
+ }
+diff --ruN a/stablehlo/stablehlo/reference/Configuration.h b/stablehlo/stablehlo/reference/Configuration.h
+--- stablehlo/stablehlo/reference/Configuration.h
++++ stablehlo/stablehlo/reference/Configuration.h
+@@ -30,7 +30,7 @@
+  public:
+   /// Custom op kernels for any user specified ops not found in the StableHLO
+   /// op dialect or StableHLO interpreter dialect.
+-  virtual llvm::Error operator()(Operation &op, Scope &scope, Process *process);
++  virtual llvm::Error operator()(Operation& op, Scope& scope, Process* process);
+ 
+   virtual ~InterpreterFallback() = default;
+ };
+diff --ruN a/stablehlo/stablehlo/reference/Errors.h b/stablehlo/stablehlo/reference/Errors.h
+--- stablehlo/stablehlo/reference/Errors.h
++++ stablehlo/stablehlo/reference/Errors.h
+@@ -26,7 +26,7 @@
+ /// error using a format string and a variadic number of arguments to the format
+ /// string.
+ template <typename... Ts>
+-inline llvm::Error invalidArgument(char const *Fmt, const Ts &...Vals) {
++inline llvm::Error invalidArgument(char const* Fmt, const Ts&... Vals) {
+   return createStringError(llvm::errc::invalid_argument, Fmt, Vals...);
+ }
+ 
+diff --ruN a/stablehlo/stablehlo/reference/Index.cpp b/stablehlo/stablehlo/reference/Index.cpp
+--- stablehlo/stablehlo/reference/Index.cpp
++++ stablehlo/stablehlo/reference/Index.cpp
+@@ -26,7 +26,7 @@
+ namespace mlir {
+ namespace stablehlo {
+ 
+-raw_ostream &operator<<(raw_ostream &os, const Sizes &x) {
++raw_ostream& operator<<(raw_ostream& os, const Sizes& x) {
+   os << "[";
+   llvm::interleave(x, os, ", ");
+   os << "]";
+@@ -40,7 +40,7 @@
+   return result;
+ }
+ 
+-bool Sizes::inBounds(const Sizes &bounds) const {
++bool Sizes::inBounds(const Sizes& bounds) const {
+   if (size() != bounds.size()) return false;
+   for (auto [size, bound] : llvm::zip(*this, bounds))
+     if (size < 0 || size >= bound) return false;
+@@ -59,7 +59,7 @@
+   return IndexSpaceIterator(*this);
+ }
+ 
+-Sizes operator+(const Sizes &x, const Sizes &y) {
++Sizes operator+(const Sizes& x, const Sizes& y) {
+   if (x.size() != y.size()) llvm::report_fatal_error("expected same size");
+   Sizes result(x.size());
+   for (size_t i = 0; i < x.size(); ++i) {
+@@ -68,11 +68,11 @@
+   return result;
+ }
+ 
+-Sizes operator+(const Sizes &x, int64_t y) { return x + Sizes(x.size(), y); }
++Sizes operator+(const Sizes& x, int64_t y) { return x + Sizes(x.size(), y); }
+ 
+-Sizes operator+(int64_t x, const Sizes &y) { return y + x; }
++Sizes operator+(int64_t x, const Sizes& y) { return y + x; }
+ 
+-Sizes operator-(const Sizes &x, const Sizes &y) {
++Sizes operator-(const Sizes& x, const Sizes& y) {
+   if (x.size() != y.size()) llvm::report_fatal_error("expected same size");
+   Sizes result(x.size());
+   for (size_t i = 0; i < x.size(); ++i) {
+@@ -81,11 +81,11 @@
+   return result;
+ }
+ 
+-Sizes operator-(const Sizes &x, int64_t y) { return x - Sizes(x.size(), y); }
++Sizes operator-(const Sizes& x, int64_t y) { return x - Sizes(x.size(), y); }
+ 
+-Sizes operator-(int64_t x, const Sizes &y) { return Sizes(y.size(), x) - y; }
++Sizes operator-(int64_t x, const Sizes& y) { return Sizes(y.size(), x) - y; }
+ 
+-Sizes operator*(const Sizes &x, const Sizes &y) {
++Sizes operator*(const Sizes& x, const Sizes& y) {
+   if (x.size() != y.size()) llvm::report_fatal_error("expected same size");
+   Sizes result(x.size());
+   for (size_t i = 0; i < x.size(); ++i) {
+@@ -94,23 +94,23 @@
+   return result;
+ }
+ 
+-Sizes operator*(const Sizes &x, int64_t y) { return x * Sizes(x.size(), y); }
++Sizes operator*(const Sizes& x, int64_t y) { return x * Sizes(x.size(), y); }
+ 
+-Sizes operator*(int64_t &x, const Sizes &y) { return y + x; }
++Sizes operator*(int64_t& x, const Sizes& y) { return y + x; }
+ 
+-Sizes clamp(int64_t min, const Sizes &x, int64_t max) {
++Sizes clamp(int64_t min, const Sizes& x, int64_t max) {
+   return clamp(Sizes(x.size(), min), x, Sizes(x.size(), max));
+ }
+ 
+-Sizes clamp(int64_t min, const Sizes &x, const Sizes &max) {
++Sizes clamp(int64_t min, const Sizes& x, const Sizes& max) {
+   return clamp(Sizes(x.size(), min), x, max);
+ }
+ 
+-Sizes clamp(const Sizes &min, const Sizes &x, int64_t max) {
++Sizes clamp(const Sizes& min, const Sizes& x, int64_t max) {
+   return clamp(min, x, Sizes(x.size(), max));
+ }
+ 
+-Sizes clamp(const Sizes &min, const Sizes &x, const Sizes &max) {
++Sizes clamp(const Sizes& min, const Sizes& x, const Sizes& max) {
+   if (min.size() != x.size() || x.size() != max.size())
+     llvm::report_fatal_error("expected same size");
+   Sizes result(x.size());
+@@ -120,15 +120,15 @@
+   return result;
+ }
+ 
+-const Index &IndexSpaceIterator::operator*() const {
++const Index& IndexSpaceIterator::operator*() const {
+   if (!index_)
+     llvm::report_fatal_error("Dereferencing a past-the-end iterator.");
+   return *index_;
+ }
+ 
+-const Index *IndexSpaceIterator::operator->() const { return &(*index_); }
++const Index* IndexSpaceIterator::operator->() const { return &(*index_); }
+ 
+-IndexSpaceIterator &IndexSpaceIterator::operator++() {
++IndexSpaceIterator& IndexSpaceIterator::operator++() {
+   if (!index_)
+     llvm::report_fatal_error("Incrementing a past-the-end iterator.");
+ 
+diff --ruN a/stablehlo/stablehlo/reference/Index.h b/stablehlo/stablehlo/reference/Index.h
+--- stablehlo/stablehlo/reference/Index.h
++++ stablehlo/stablehlo/reference/Index.h
+@@ -32,8 +32,8 @@
+ class Sizes : public SmallVector<int64_t> {
+  public:
+   Sizes() = default;
+-  Sizes(const Sizes &other) = default;
+-  Sizes &operator=(const Sizes &other) = default;
++  Sizes(const Sizes& other) = default;
++  Sizes& operator=(const Sizes& other) = default;
+ 
+   Sizes(std::initializer_list<int64_t> list) : SmallVector(list) {}
+   Sizes(iterator begin, iterator end) : SmallVector(begin, end) {}
+@@ -49,66 +49,66 @@
+ 
+   /// Checks if an element `e` at kth axis of `this` object follows
+   /// `0 <= e <= bounds[k]`.
+-  bool inBounds(const Sizes &bounds) const;
++  bool inBounds(const Sizes& bounds) const;
+ 
+   /// Iterate over the index space of a Sizes object.
+   IndexSpaceIterator index_begin() const;
+   IndexSpaceIterator index_end() const;
+ };
+ 
+-raw_ostream &operator<<(raw_ostream &os, const Sizes &x);
++raw_ostream& operator<<(raw_ostream& os, const Sizes& x);
+ 
+ /// Overloaded add operator to return `Sizes` object `z` such that
+ /// `z[k] = x[k] + y[k]` for all axis k.
+-Sizes operator+(const Sizes &x, const Sizes &y);
++Sizes operator+(const Sizes& x, const Sizes& y);
+ 
+ /// Overloaded add operator to return `Sizes` object `z` such that
+ /// `z[k] = x[k] + y` for all axis k.
+-Sizes operator+(const Sizes &x, int64_t y);
++Sizes operator+(const Sizes& x, int64_t y);
+ 
+ /// Overloaded add operator to return `Sizes` object `z` such that
+ /// `z[k] = x + y[k]` for all axis k.
+-Sizes operator+(int64_t x, const Sizes &y);
++Sizes operator+(int64_t x, const Sizes& y);
+ 
+ /// Overloaded add operator to return `Sizes` object `z` such that
+ /// `z[k] = x[k] - y[k]` for all axis k.
+-Sizes operator-(const Sizes &x, const Sizes &y);
++Sizes operator-(const Sizes& x, const Sizes& y);
+ 
+ /// Overloaded add operator to return `Sizes` object `z` such that
+ /// `z[k] = x[k] - y` for all axis k.
+-Sizes operator-(const Sizes &x, int64_t y);
++Sizes operator-(const Sizes& x, int64_t y);
+ 
+ /// Overloaded add operator to return `Sizes` object `z` such that
+ /// `z[k] = x - y[k]` for all axis k.
+-Sizes operator-(int64_t x, const Sizes &y);
++Sizes operator-(int64_t x, const Sizes& y);
+ 
+ /// Overloaded add operator to return `Sizes` object `z` such that
+ /// `z[k] = x[k] * y[k]` for all axis k.
+-Sizes operator*(const Sizes &x, const Sizes &y);
++Sizes operator*(const Sizes& x, const Sizes& y);
+ 
+ /// Overloaded add operator to return `Sizes` object `z` such that
+ /// `z[k] = x[k] * y` for all axis k.
+-Sizes operator*(const Sizes &x, int64_t y);
++Sizes operator*(const Sizes& x, int64_t y);
+ 
+ /// Overloaded add operator to return `Sizes` object `z` such that
+ /// `z[k] = x * y[k]` for all axis k.
+-Sizes operator*(int64_t x, const Sizes &y);
++Sizes operator*(int64_t x, const Sizes& y);
+ 
+ /// Clamp operator to return `Sizes` object `z` such that
+ /// `z[k] = std::min(std::max(x[k], min[k]), max[k])` for all axis k.
+-Sizes clamp(const Sizes &min, const Sizes &x, const Sizes &max);
++Sizes clamp(const Sizes& min, const Sizes& x, const Sizes& max);
+ 
+ /// Clamp operator to return `Sizes` object `z` such that
+ /// `z[k] = std::min(std::max(x[k], min), max)` for all axis k.
+-Sizes clamp(int64_t min, const Sizes &x, int64_t max);
++Sizes clamp(int64_t min, const Sizes& x, int64_t max);
+ 
+ /// Clamp operator to return `Sizes` object `z` such that
+ /// `z[k] = std::min(std::max(x[k], min), max[k])` for all axis k.
+-Sizes clamp(int64_t min, const Sizes &x, const Sizes &max);
++Sizes clamp(int64_t min, const Sizes& x, const Sizes& max);
+ 
+ /// Clamp operator to return `Sizes` object `z` such that
+ /// `z[k] = std::min(std::max(x[k], min[k]), max)` for all axis k.
+-Sizes clamp(const Sizes &min, const Sizes &x, int64_t max);
++Sizes clamp(const Sizes& min, const Sizes& x, int64_t max);
+ 
+ /// Represents index of a tensor.
+ using Index = Sizes;
+@@ -132,23 +132,23 @@
+   /// At any point in time, the iterator can either reference an actual index
+   /// or the past-the-end element in the index space.
+   /// Dereferencing a past-the-end iterator will result in a fatal error.
+-  const Index &operator*() const;
+-  const Index *operator->() const;
++  const Index& operator*() const;
++  const Index* operator->() const;
+ 
+   /// Compare the iterator to another iterator.
+   /// Two iterators are equal if they have the same underlying shape and
+   /// reference the same element in the index space.
+-  bool operator==(const IndexSpaceIterator &it) {
++  bool operator==(const IndexSpaceIterator& it) {
+     return shape_ == it.shape_ && index_ == it.index_;
+   }
+-  bool operator!=(const IndexSpaceIterator &it) { return !(*this == it); }
++  bool operator!=(const IndexSpaceIterator& it) { return !(*this == it); }
+ 
+   /// Increment to the next index while iterating over the index space
+   /// of a tensor in lexicographical order.
+   /// Incrementing past the last index will result in a past-the-end iterator
+   /// which cannot be dereferenced. Incrementing even further will result in
+   /// a fatal error.
+-  IndexSpaceIterator &operator++();
++  IndexSpaceIterator& operator++();
+   IndexSpaceIterator operator++(int);
+ 
+  private:
+diff --ruN a/stablehlo/stablehlo/reference/Process.cpp b/stablehlo/stablehlo/reference/Process.cpp
+--- stablehlo/stablehlo/reference/Process.cpp
++++ stablehlo/stablehlo/reference/Process.cpp
+@@ -25,7 +25,7 @@
+ namespace mlir {
+ namespace stablehlo {
+ 
+-Process::Process(ProcessId id, ProcessGrid *grid) : id_(id), grid_(grid) {}
++Process::Process(ProcessId id, ProcessGrid* grid) : id_(id), grid_(grid) {}
+ 
+ ProcessGroups Process::crossPartition(
+     SmallVector<SmallVector<uint32_t>> partitionGroups) {
+diff --ruN a/stablehlo/stablehlo/reference/Process.h b/stablehlo/stablehlo/reference/Process.h
+--- stablehlo/stablehlo/reference/Process.h
++++ stablehlo/stablehlo/reference/Process.h
+@@ -29,7 +29,7 @@
+  public:
+   /// \name Constructors
+   /// @{
+-  Process(ProcessId id, ProcessGrid *grid);
++  Process(ProcessId id, ProcessGrid* grid);
+   /// @}
+ 
+   /// See `ProcessGrid::crossPartition`.
+@@ -72,7 +72,7 @@
+ 
+   /// See ProcessGrid. The pointer is used to gain access to allow
+   /// synchronization among participating processes.
+-  ProcessGrid *grid_;
++  ProcessGrid* grid_;
+ };
+ 
+ }  // namespace stablehlo
+diff --ruN a/stablehlo/stablehlo/reference/ProcessGrid.cpp b/stablehlo/stablehlo/reference/ProcessGrid.cpp
+--- stablehlo/stablehlo/reference/ProcessGrid.cpp
++++ stablehlo/stablehlo/reference/ProcessGrid.cpp
+@@ -39,16 +39,16 @@
+ // ProcessId.
+ //===----------------------------------------------------------------------===//
+ 
+-bool ProcessId::operator!=(const ProcessId &other) const {
++bool ProcessId::operator!=(const ProcessId& other) const {
+   return !(*this == other);
+ }
+ 
+-bool ProcessId::operator<(const ProcessId &other) const {
++bool ProcessId::operator<(const ProcessId& other) const {
+   return std::pair<uint32_t, uint32_t>{replicaId, partitionId} <
+          std::pair<uint32_t, uint32_t>{other.replicaId, other.partitionId};
+ }
+ 
+-bool ProcessId::operator==(const ProcessId &other) const {
++bool ProcessId::operator==(const ProcessId& other) const {
+   return std::pair<uint32_t, uint32_t>{replicaId, partitionId} ==
+          std::pair<uint32_t, uint32_t>{other.replicaId, other.partitionId};
+ }
+@@ -70,7 +70,7 @@
+ //===----------------------------------------------------------------------===//
+ 
+ RendezvousResult::RendezvousResult(
+-    std::map<ProcessId, SmallVector<Tensor>> const &results)
++    std::map<ProcessId, SmallVector<Tensor>> const& results)
+     : results_(results) {}
+ 
+ void RendezvousResult::insert(ProcessId processId,
+@@ -86,12 +86,12 @@
+ 
+ SmallVector<SmallVector<Tensor>> RendezvousResult::getSortedTensors() const {
+   return llvm::map_to_vector(results_,
+-                             [](const auto &pair) { return pair.second; });
++                             [](const auto& pair) { return pair.second; });
+ }
+ 
+ bool RendezvousResult::hasMatchingOperandsCount() const {
+   auto count = results_.begin()->second.size();
+-  for (const auto &it : results_)
++  for (const auto& it : results_)
+     if (count != it.second.size()) return false;
+   return true;
+ }
+@@ -101,7 +101,7 @@
+ //===----------------------------------------------------------------------===//
+ 
+ template <typename K, typename V>
+-V &detail::ThreadSafeMap<K, V>::operator[](const K &key) {
++V& detail::ThreadSafeMap<K, V>::operator[](const K& key) {
+   std::lock_guard<std::mutex> lock(lock_);
+   return map_[key];
+ }
+@@ -133,7 +133,7 @@
+ //===----------------------------------------------------------------------===//
+ 
+ template <typename T>
+-detail::ThreadSafeQueue<T>::ThreadSafeQueue(const std::queue<T> &queue)
++detail::ThreadSafeQueue<T>::ThreadSafeQueue(const std::queue<T>& queue)
+     : queue_(queue) {}
+ 
+ template <typename T>
+@@ -155,7 +155,7 @@
+ //===----------------------------------------------------------------------===//
+ 
+ ProcessGrid::ProcessGrid(uint32_t numReplicas, uint32_t numPartitions,
+-                         std::queue<StringAttr> &infeed)
++                         std::queue<StringAttr>& infeed)
+     : numReplicas_(numReplicas),
+       numPartitions_(numPartitions),
+       infeed_(infeed) {}
+@@ -163,7 +163,7 @@
+ ProcessGroups ProcessGrid::crossPartition(
+     SmallVector<SmallVector<uint32_t>> partitionGroups) {
+   ProcessGroups processGroups;
+-  for (const auto &partitionGroup : partitionGroups) {
++  for (const auto& partitionGroup : partitionGroups) {
+     for (uint32_t replicaId = 0; replicaId < numReplicas_; ++replicaId) {
+       ProcessGroup processGroup;
+       for (uint32_t partitionId : partitionGroup)
+@@ -177,7 +177,7 @@
+ ProcessGroups ProcessGrid::crossReplica(
+     SmallVector<SmallVector<uint32_t>> replicaGroups) {
+   ProcessGroups processGroups;
+-  for (const auto &replicaGroup : replicaGroups) {
++  for (const auto& replicaGroup : replicaGroups) {
+     for (uint32_t partitionId = 0; partitionId < numPartitions_;
+          ++partitionId) {
+       ProcessGroup processGroup;
+@@ -192,7 +192,7 @@
+ ProcessGroups ProcessGrid::crossReplicaAndPartition(
+     SmallVector<SmallVector<uint32_t>> replicaGroups) {
+   ProcessGroups processGroups;
+-  for (const auto &replicaGroup : replicaGroups) {
++  for (const auto& replicaGroup : replicaGroups) {
+     ProcessGroup processGroup;
+     for (uint32_t partitionId = 0; partitionId < numPartitions_; ++partitionId)
+       for (uint32_t replicaId : replicaGroup)
+@@ -205,7 +205,7 @@
+ ProcessGroups ProcessGrid::flattenedIds(
+     SmallVector<SmallVector<uint32_t>> flattenedIdGroups) {
+   ProcessGroups processGroups;
+-  for (const auto &flattenedIdGroup : flattenedIdGroups) {
++  for (const auto& flattenedIdGroup : flattenedIdGroups) {
+     ProcessGroup processGroup;
+     for (auto flattenedId : flattenedIdGroup) {
+       uint32_t replicaId = flattenedId / numPartitions_;
+@@ -251,7 +251,7 @@
+   }
+ 
+   std::pair<ProcessGroup, ChannelId> channelKey(processGroup, channelId);
+-  auto &state = channels_[channelKey];
++  auto& state = channels_[channelKey];
+ 
+   std::unique_lock<std::mutex> lock(state.mutex);
+   state.values[processId] = SmallVector<Tensor>(operands);
+diff --ruN a/stablehlo/stablehlo/reference/ProcessGrid.h b/stablehlo/stablehlo/reference/ProcessGrid.h
+--- stablehlo/stablehlo/reference/ProcessGrid.h
++++ stablehlo/stablehlo/reference/ProcessGrid.h
+@@ -41,7 +41,7 @@
+ class RendezvousResult {
+  public:
+   RendezvousResult() = default;
+-  RendezvousResult(std::map<ProcessId, SmallVector<Tensor>> const &results);
++  RendezvousResult(std::map<ProcessId, SmallVector<Tensor>> const& results);
+ 
+   /// Iterates through the (ProcessId, SmallVector<Tensor>) map entires and
+   /// returns a vector of Tensors sorted by ProcessId--(replicaId, partitionId)
+@@ -104,7 +104,7 @@
+ class ThreadSafeMap {
+  public:
+   /// Returns a reference to the data associated with the `key`.
+-  V &operator[](const K &key);
++  V& operator[](const K& key);
+ 
+  private:
+   /// Synchronization primitive used to manage concurrent access to the map.
+@@ -143,7 +143,7 @@
+   /// \name Constructors
+   /// @{
+   ThreadSafeQueue() = default;
+-  ThreadSafeQueue(const std::queue<T> &queue);
++  ThreadSafeQueue(const std::queue<T>& queue);
+   /// @}
+ 
+   /// Remove the first element of the queue and return it.
+@@ -173,16 +173,16 @@
+   uint32_t partitionId;
+ 
+   /// Overloaded inequality operator.
+-  bool operator!=(const ProcessId &other) const;
++  bool operator!=(const ProcessId& other) const;
+ 
+   /// The sort order for ProcessId is not defined in StableHLO, and it's
+   /// internally used in ProcessGrid::rendezvous as part of a sorted key on the
+   /// map. This operator is conveniently used to help define the ordering since
+   /// ordering is defined for StableHLO process group.
+-  bool operator<(const ProcessId &other) const;
++  bool operator<(const ProcessId& other) const;
+ 
+   /// Overloaded equality operator.
+-  bool operator==(const ProcessId &other) const;
++  bool operator==(const ProcessId& other) const;
+ };
+ 
+ /// StableHLO `process_group`.
+@@ -190,9 +190,9 @@
+  public:
+   using SmallVector<ProcessId>::SmallVector;
+ 
+-  bool operator<(const ProcessGroup &other) const {
+-    return static_cast<const SmallVector &>(*this) <
+-           static_cast<const SmallVector &>(other);
++  bool operator<(const ProcessGroup& other) const {
++    return static_cast<const SmallVector&>(*this) <
++           static_cast<const SmallVector&>(other);
+   }
+ };
+ 
+@@ -211,7 +211,7 @@
+   /// \name Constructors
+   /// @{
+   ProcessGrid(uint32_t numReplicas, uint32_t numPartitions,
+-              std::queue<StringAttr> &infeed);
++              std::queue<StringAttr>& infeed);
+   /// @}
+ 
+   /// StableHLO `cross_partition` communication strategy.
+diff --ruN a/stablehlo/stablehlo/reference/Token.cpp b/stablehlo/stablehlo/reference/Token.cpp
+--- stablehlo/stablehlo/reference/Token.cpp
++++ stablehlo/stablehlo/reference/Token.cpp
+@@ -23,11 +23,11 @@
+ namespace mlir {
+ namespace stablehlo {
+ 
+-Token::Token(MLIRContext *context) : type_(TokenType::get(context)) {}
++Token::Token(MLIRContext* context) : type_(TokenType::get(context)) {}
+ 
+ TokenType Token::getType() const { return type_; }
+ 
+-void Token::print(raw_ostream &os) const { getType().print(os); }
++void Token::print(raw_ostream& os) const { getType().print(os); }
+ 
+ void Token::dump() const { print(llvm::errs()); }
+ 
+diff --ruN a/stablehlo/stablehlo/reference/Token.h b/stablehlo/stablehlo/reference/Token.h
+--- stablehlo/stablehlo/reference/Token.h
++++ stablehlo/stablehlo/reference/Token.h
+@@ -28,14 +28,14 @@
+  public:
+   /// \name Constructors
+   /// @{
+-  Token(MLIRContext *context);
++  Token(MLIRContext* context);
+   /// @}
+ 
+   /// Returns the type of the Token object.
+   TokenType getType() const;
+ 
+   /// Prints Token object.
+-  void print(raw_ostream &os) const;
++  void print(raw_ostream& os) const;
+   void dump() const;
+ 
+  private:
+@@ -43,7 +43,7 @@
+ };
+ 
+ /// Print utilities for Token objects.
+-inline raw_ostream &operator<<(raw_ostream &os, Token Token) {
++inline raw_ostream& operator<<(raw_ostream& os, Token Token) {
+   Token.print(os);
+   return os;
+ }
+diff --ruN a/stablehlo/stablehlo/reference/Value.cpp b/stablehlo/stablehlo/reference/Value.cpp
+--- stablehlo/stablehlo/reference/Value.cpp
++++ stablehlo/stablehlo/reference/Value.cpp
+@@ -46,7 +46,7 @@
+ 
+ TupleType Tuple::getType() const { return type_; }
+ 
+-void Tuple::print(raw_ostream &os) const {
++void Tuple::print(raw_ostream& os) const {
+   getType().print(os);
+   os << " (\n";
+   for (size_t i = 0; i < values_.size(); ++i) {
+@@ -63,9 +63,9 @@
+ // InterpreterValue.
+ //===----------------------------------------------------------------------===//
+ 
+-InterpreterValue::InterpreterValue(const Tensor &tensor) : value_(tensor) {}
+-InterpreterValue::InterpreterValue(const Token &token) : value_(token) {}
+-InterpreterValue::InterpreterValue(const Tuple &tuple) : value_(tuple) {}
++InterpreterValue::InterpreterValue(const Tensor& tensor) : value_(tensor) {}
++InterpreterValue::InterpreterValue(const Token& token) : value_(token) {}
++InterpreterValue::InterpreterValue(const Tuple& tuple) : value_(tuple) {}
+ 
+ Tensor InterpreterValue::getTensor() const {
+   if (!isTensor())
+@@ -108,7 +108,7 @@
+   return std::holds_alternative<Tuple>(value_);
+ }
+ 
+-void InterpreterValue::print(raw_ostream &os) const {
++void InterpreterValue::print(raw_ostream& os) const {
+   if (isTensor())
+     getTensor().print(os);
+   else if (isToken())
 diff --ruN a/stablehlo/stablehlo/testdata/BUILD.bazel b/stablehlo/stablehlo/testdata/BUILD.bazel
 --- stablehlo/stablehlo/testdata/BUILD.bazel
 +++ stablehlo/stablehlo/testdata/BUILD.bazel
-@@ -11,13 +11,26 @@
+@@ -11,14 +11,27 @@
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
@@ -2832,7 +3948,7 @@ diff --ruN a/stablehlo/stablehlo/testdata/BUILD.bazel b/stablehlo/stablehlo/test
 +    ],
      licenses = ["notice"],
  )
-+
+ 
 +# copybara:uncomment_begin
 +# filegroup(
 +#     name = "all_files",
@@ -2840,9 +3956,10 @@ diff --ruN a/stablehlo/stablehlo/testdata/BUILD.bazel b/stablehlo/stablehlo/test
 +#     visibility = ["//:__subpackages__"],
 +# )
 +# copybara:uncomment_end
- 
++
  # Equivalent of configure_lit_site_cfg from CMakeLists.txt.
  expand_template(
+     name = "lit_site_cfg_py_gen",
 @@ -39,46 +52,42 @@
      template = "lit.site.cfg.py.in",
  )
@@ -2926,7 +4043,7 @@ diff --ruN a/stablehlo/stablehlo/testdata/BUILD.bazel b/stablehlo/stablehlo/test
 diff --ruN a/stablehlo/stablehlo/tests/BUILD.bazel b/stablehlo/stablehlo/tests/BUILD.bazel
 --- stablehlo/stablehlo/tests/BUILD.bazel
 +++ stablehlo/stablehlo/tests/BUILD.bazel
-@@ -11,15 +11,29 @@
+@@ -11,16 +11,30 @@
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
@@ -2947,7 +4064,7 @@ diff --ruN a/stablehlo/stablehlo/tests/BUILD.bazel b/stablehlo/stablehlo/tests/B
 +    ],
      licenses = ["notice"],
  )
-+
+ 
 +# copybara:uncomment_begin
 +# filegroup(
 +#     name = "all_files",
@@ -2955,9 +4072,10 @@ diff --ruN a/stablehlo/stablehlo/tests/BUILD.bazel b/stablehlo/stablehlo/tests/B
 +#     visibility = ["//:__subpackages__"],
 +# )
 +# copybara:uncomment_end
- 
++
  cc_library(
      name = "check_ops",
+     srcs = [
 @@ -29,6 +43,7 @@
      hdrs = [
          "CheckOps.h",
@@ -2994,13 +4112,12 @@ diff --ruN a/stablehlo/stablehlo/tests/BUILD.bazel b/stablehlo/stablehlo/tests/B
  
  gentbl_cc_library(
      name = "test_utils_inc_gen",
--    strip_include_prefix = ".",
++    compatible_with = get_compatible_with_portable(),
+     strip_include_prefix = ".",
 -    tbl_outs = {"TestUtils.h.inc": [
 -        "-gen-pass-decls",
 -        "-name=HloTest",
 -    ]},
-+    compatible_with = get_compatible_with_portable(),
-+    strip_include_prefix = ".",
 +    tbl_outs = {
 +        "TestUtils.h.inc": [
 +            "-gen-pass-decls",
@@ -3045,8 +4162,10 @@ diff --ruN a/stablehlo/stablehlo/tests/BUILD.bazel b/stablehlo/stablehlo/tests/B
 -    )
 -    for src in glob(["**/*.mlir"])
 -]
+-
+-test_suite(
 +lit_test_suite(
-+    name = "stablehlo_tests",
+     name = "stablehlo_tests",
 +    srcs = glob([
 +        "*.mlir",
 +        "**/*.mlir",
@@ -3060,18 +4179,116 @@ diff --ruN a/stablehlo/stablehlo/tests/BUILD.bazel b/stablehlo/stablehlo/tests/B
 +        "@llvm-project//llvm:not",
 +    ],
 +    size_per_test = "small",
-+    tags = ["stablehlo_tests"],
-+)
- 
- test_suite(
--    name = "stablehlo_tests",
--    tags = ["stablehlo_tests"],
--)
+     tags = ["stablehlo_tests"],
+ )
++
++test_suite(
 +    name = "all_tests",  # alias for `:stablehlo_tests`
 +    tests = [
 +        "//stablehlo/tests:stablehlo_tests",
 +    ],
 +)
+diff --ruN a/stablehlo/stablehlo/tests/CheckOps.cpp b/stablehlo/stablehlo/tests/CheckOps.cpp
+--- stablehlo/stablehlo/tests/CheckOps.cpp
++++ stablehlo/stablehlo/tests/CheckOps.cpp
+@@ -72,7 +72,7 @@
+ // Check Dialect Constructor
+ //===----------------------------------------------------------------------===//
+ 
+-CheckDialect::CheckDialect(MLIRContext *context)
++CheckDialect::CheckDialect(MLIRContext* context)
+     : Dialect(getDialectNamespace(), context, TypeID::get<CheckDialect>()) {
+   addOperations<
+ #define GET_OP_LIST
+@@ -80,13 +80,13 @@
+       >();
+ }
+ 
+-llvm::Error evalExpectAlmostEqConstOp(const Tensor &lhs, ElementsAttr value,
++llvm::Error evalExpectAlmostEqConstOp(const Tensor& lhs, ElementsAttr value,
+                                       APFloat tolerance) {
+   auto rhs = makeTensor(cast<DenseElementsAttr>(value));
+   return evalExpectAlmostEqOp(lhs, rhs, tolerance);
+ }
+ 
+-llvm::Error evalExpectAlmostEqOp(const Tensor &lhs, const Tensor &rhs,
++llvm::Error evalExpectAlmostEqOp(const Tensor& lhs, const Tensor& rhs,
+                                  APFloat tolerance) {
+   for (auto lhsIt = lhs.index_begin(), rhsIt = rhs.index_begin();
+        lhsIt != lhs.index_end(); ++lhsIt, ++rhsIt)
+@@ -101,12 +101,12 @@
+   return llvm::Error::success();
+ }
+ 
+-llvm::Error evalExpectEqConstOp(const Tensor &lhs, ElementsAttr value) {
++llvm::Error evalExpectEqConstOp(const Tensor& lhs, ElementsAttr value) {
+   auto rhs = makeTensor(cast<DenseElementsAttr>(value));
+   return evalExpectEqOp(lhs, rhs);
+ }
+ 
+-llvm::Error evalExpectEqOp(const Tensor &lhs, const Tensor &rhs) {
++llvm::Error evalExpectEqOp(const Tensor& lhs, const Tensor& rhs) {
+   for (auto lhsIt = lhs.index_begin(), rhsIt = rhs.index_begin();
+        lhsIt != lhs.index_end(); ++lhsIt, ++rhsIt)
+     if ((lhs.get(*lhsIt) != rhs.get(*rhsIt)).getBooleanValue())
+@@ -146,7 +146,7 @@
+   return llvm::errc::bad_address;
+ }
+ 
+-llvm::Error evalExpectSerializedEqOp(const Tensor &expected, StringRef probeId,
++llvm::Error evalExpectSerializedEqOp(const Tensor& expected, StringRef probeId,
+                                      StringRef probeDir, uint32_t iteration) {
+   auto serializedMetadata =
+       getSerializedTensorMetadata(probeId, probeDir, iteration);
+@@ -206,7 +206,7 @@
+   return std::numeric_limits<uint64_t>::max();
+ }
+ 
+-static uint64_t ULPDifference(const Element &e1, const Element &e2) {
++static uint64_t ULPDifference(const Element& e1, const Element& e2) {
+   // caller is responsible for ensuring that e1, e2 have both the same
+   // float or complex types
+   if (isSupportedComplexType(e1.getType())) {
+@@ -218,7 +218,7 @@
+   return ULPDifference(e1.getFloatValue(), e2.getFloatValue());
+ }
+ 
+-llvm::Error evalExpectCloseOp(const Tensor &actual, const Tensor &expected,
++llvm::Error evalExpectCloseOp(const Tensor& actual, const Tensor& expected,
+                               uint64_t min_ulp_difference,
+                               uint64_t max_ulp_difference) {
+   auto type = actual.getElementType();
+diff --ruN a/stablehlo/stablehlo/tests/CheckOps.h b/stablehlo/stablehlo/tests/CheckOps.h
+--- stablehlo/stablehlo/tests/CheckOps.h
++++ stablehlo/stablehlo/tests/CheckOps.h
+@@ -32,20 +32,20 @@
+ 
+ class CheckDialect : public Dialect {
+  public:
+-  explicit CheckDialect(MLIRContext *context);
++  explicit CheckDialect(MLIRContext* context);
+   static StringRef getDialectNamespace() { return "check"; }
+ };
+ 
+ // The eval functions for the following ops are used only for test harness.
+-llvm::Error evalExpectAlmostEqConstOp(const Tensor &lhs, ElementsAttr value,
++llvm::Error evalExpectAlmostEqConstOp(const Tensor& lhs, ElementsAttr value,
+                                       APFloat tolerance);
+-llvm::Error evalExpectAlmostEqOp(const Tensor &lhs, const Tensor &rhs,
++llvm::Error evalExpectAlmostEqOp(const Tensor& lhs, const Tensor& rhs,
+                                  APFloat tolerance);
+-llvm::Error evalExpectEqConstOp(const Tensor &lhs, ElementsAttr value);
+-llvm::Error evalExpectEqOp(const Tensor &lhs, const Tensor &rhs);
+-llvm::Error evalExpectSerializedEqOp(const Tensor &expected, StringRef probeId,
++llvm::Error evalExpectEqConstOp(const Tensor& lhs, ElementsAttr value);
++llvm::Error evalExpectEqOp(const Tensor& lhs, const Tensor& rhs);
++llvm::Error evalExpectSerializedEqOp(const Tensor& expected, StringRef probeId,
+                                      StringRef probeDir, uint32_t iteration);
+-llvm::Error evalExpectCloseOp(const Tensor &actual, const Tensor &expected,
++llvm::Error evalExpectCloseOp(const Tensor& actual, const Tensor& expected,
+                               uint64_t min_ulp_difference,
+                               uint64_t max_ulp_difference);
+ 
 diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir
 --- stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir
 +++ stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir
@@ -3324,14 +4541,10 @@ diff --ruN a/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir b/stablehlo
 diff --ruN a/stablehlo/stablehlo/tests/ops_chlo.mlir b/stablehlo/stablehlo/tests/ops_chlo.mlir
 --- stablehlo/stablehlo/tests/ops_chlo.mlir
 +++ stablehlo/stablehlo/tests/ops_chlo.mlir
-@@ -485,6 +485,50 @@
-     stablehlo.return %1, %2, %1, %2 : tensor<3xf32>, tensor<3xi32>, tensor<3xf32>, tensor<3xi32>
-   } : (tensor<2x3xf32>, tensor<2x3xi32>, tensor<3xf32>, tensor<3xi32>) -> (tensor<2x3xf32>, tensor<2x3xi32>, tensor<3xf32>, tensor<3xi32>)
-   func.return %0#0, %0#1 : tensor<2x3xf32>, tensor<2x3xi32>
-+}
-+
-+// -----
-+
+@@ -489,6 +489,50 @@
+ 
+ // -----
+ 
 +// CHECK-LABEL: func @scan_bounded_dynamism
 +func.func @scan_bounded_dynamism(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, %arg1: tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>) -> (tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>) {
 +  // CHECK: chlo.scan
@@ -3372,9 +4585,13 @@ diff --ruN a/stablehlo/stablehlo/tests/ops_chlo.mlir b/stablehlo/stablehlo/tests
 +    stablehlo.return %1, %2, %1, %2 : tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?xi32, #stablehlo.type_extensions<bounds = [4]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?xi32, #stablehlo.type_extensions<bounds = [4]>>
 +  } : (tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?x?xi32, #stablehlo.type_extensions<bounds = [2, 4]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?xi32, #stablehlo.type_extensions<bounds = [4]>>) -> (tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?x?xi32, #stablehlo.type_extensions<bounds = [2, 4]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?xi32, #stablehlo.type_extensions<bounds = [4]>>)
 +  func.return %0#0, %0#1, %0#2, %0#3 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?x?xi32, #stablehlo.type_extensions<bounds = [2, 4]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?xi32, #stablehlo.type_extensions<bounds = [4]>>
- }
- 
- // -----
++}
++
++// -----
++
+ func.func @scan_explicit_size_mismatch(%arg0: tensor<2x3xf32>, %arg1: tensor<3xf32>) -> tensor<2x3xf32> {
+   // expected-error @+1 {{'chlo.scan' op invalid scan dimension size of operand 0}}
+   %0, %1 = chlo.scan (%arg0) inits (%arg1) dimension = 0 attributes { scan_dim_size = 3 : i64 } {
 @@ -561,6 +605,30 @@
  
  // -----
@@ -3561,13 +4778,181 @@ diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir b
      %2 = stablehlo.dynamic_iota %1, dim = 0 : (tensor<1xi64>) -> tensor<?xi32>
      return %2 : tensor<?xi32>
    }
+diff --ruN a/stablehlo/stablehlo/tools/StablehloLspServerMain.cpp b/stablehlo/stablehlo/tools/StablehloLspServerMain.cpp
+--- stablehlo/stablehlo/tools/StablehloLspServerMain.cpp
++++ stablehlo/stablehlo/tools/StablehloLspServerMain.cpp
+@@ -19,7 +19,7 @@
+ #include "mlir/Tools/mlir-lsp-server/MlirLspServerMain.h"
+ #include "stablehlo/dialect/Register.h"
+ 
+-int main(int argc, char **argv) {
++int main(int argc, char** argv) {
+   mlir::DialectRegistry registry;
+   mlir::registerAllDialects(registry);
+   mlir::registerAllExtensions(registry);
+diff --ruN a/stablehlo/stablehlo/tools/StablehloOptMain.cpp b/stablehlo/stablehlo/tools/StablehloOptMain.cpp
+--- stablehlo/stablehlo/tools/StablehloOptMain.cpp
++++ stablehlo/stablehlo/tools/StablehloOptMain.cpp
+@@ -28,7 +28,7 @@
+ #include "stablehlo/transforms/Passes.h"
+ #include "stablehlo/transforms/optimization/Passes.h"
+ 
+-int main(int argc, char **argv) {
++int main(int argc, char** argv) {
+   mlir::registerAllPasses();
+   mlir::hlo::registerAllTestPasses();
+   mlir::stablehlo::registerPassPipelines();
+diff --ruN a/stablehlo/stablehlo/tools/StablehloTranslateMain.cpp b/stablehlo/stablehlo/tools/StablehloTranslateMain.cpp
+--- stablehlo/stablehlo/tools/StablehloTranslateMain.cpp
++++ stablehlo/stablehlo/tools/StablehloTranslateMain.cpp
+@@ -92,7 +92,7 @@
+ 
+ namespace {
+ 
+-stablehlo::Tensor makeBooleanTensor(MLIRContext *context, bool value) {
++stablehlo::Tensor makeBooleanTensor(MLIRContext* context, bool value) {
+   auto builder = Builder(context);
+   auto type = RankedTensorType::get({}, builder.getI1Type());
+   auto res = DenseElementsAttr::get(type, builder.getBoolAttr(true));
+@@ -104,7 +104,7 @@
+ //   --args=[dense<1> : tensor<2xi32>, ...], where each dense attribute is
+ //   interpreted as a tensor.
+ mlir::FailureOr<SmallVector<stablehlo::InterpreterValue>>
+-parseInterpreterArguments(std::string argsStr, MLIRContext *context) {
++parseInterpreterArguments(std::string argsStr, MLIRContext* context) {
+   llvm::SmallVector<stablehlo::InterpreterValue> inputs;
+   auto parseError = [&](llvm::StringRef msg) {
+     std::string usage = "--args=[dense<1> : tensor<2xi32>, ...]";
+@@ -129,7 +129,7 @@
+ }
+ 
+ llvm::Error evalCustomCallCheckEq(stablehlo::CustomCallOp op,
+-                                  stablehlo::Scope &scope) {
++                                  stablehlo::Scope& scope) {
+   if (op->getNumOperands() != 2)
+     return stablehlo::invalidArgument("Unsupported op: %s",
+                                       debugString(op).c_str());
+@@ -157,10 +157,10 @@
+     : public stablehlo::InterpreterFallback {
+  public:
+   StablehloTranslateInterpreterFallback(
+-      const std::string &probeInstrumentationDir)
++      const std::string& probeInstrumentationDir)
+       : probeInstrumentationDir(probeInstrumentationDir) {}
+-  virtual llvm::Error operator()(Operation &op, stablehlo::Scope &scope,
+-                                 stablehlo::Process *process) final {
++  virtual llvm::Error operator()(Operation& op, stablehlo::Scope& scope,
++                                 stablehlo::Process* process) final {
+     llvm::StringRef funcName = op.getParentOfType<func::FuncOp>().getSymName();
+     if (auto customCall = dyn_cast<stablehlo::CustomCallOp>(op)) {
+       auto callTarget = customCall.getCallTargetName();
+@@ -271,7 +271,7 @@
+ 
+ TranslateFromMLIRRegistration interpretRegistration(
+     "interpret", "Interpreter for StableHLO",
+-    [](ModuleOp module, raw_ostream &os) -> LogicalResult {
++    [](ModuleOp module, raw_ostream& os) -> LogicalResult {
+       stablehlo::InterpreterConfiguration config;
+       config.probeInstrumentationDir = probeOutputDir.getValue();
+       config.fallback = std::make_unique<StablehloTranslateInterpreterFallback>(
+@@ -286,21 +286,21 @@
+ 
+       if (interpreterPrintDense.getValue()) {
+         DenseElementsAttr denseResult;
+-        for (auto &result : *results) {
++        for (auto& result : *results) {
+           denseResult = makeDenseElementsAttr(result.getTensor());
+           os << denseResult << '\n';
+         }
+         return success();
+       }
+ 
+-      for (auto &result : *results) {
++      for (auto& result : *results) {
+         result.print(os);
+         os << '\n';
+       }
+ 
+       return success();
+     },
+-    [](DialectRegistry &registry) {
++    [](DialectRegistry& registry) {
+       registry.insert<func::FuncDialect>();
+       registry.insert<quant::QuantDialect>();
+       registry.insert<stablehlo::check::CheckDialect>();
+@@ -310,7 +310,7 @@
+ 
+ TranslateFromMLIRRegistration serializeRegistration(
+     "serialize", "Serialize StableHLO program into a portable artifact",
+-    [](ModuleOp module, raw_ostream &os) -> LogicalResult {
++    [](ModuleOp module, raw_ostream& os) -> LogicalResult {
+       std::string targetVersion = targetOption.getValue();
+       if (targetVersion == "current")
+         targetVersion = vhlo::Version::getCurrentVersion().toString();
+@@ -325,14 +325,14 @@
+       return stablehlo::serializePortableArtifact(
+           module, targetVersion, os, allowOtherDialectsOption.getValue());
+     },
+-    [](DialectRegistry &registry) {
++    [](DialectRegistry& registry) {
+       mlir::registerAllDialects(registry);
+       mlir::stablehlo::registerAllDialects(registry);
+     });
+ 
+ TranslateToMLIRRegistration deserializeRegistration(
+     "deserialize", "Deserialize a portable artifact into a StableHLO program",
+-    [](llvm::StringRef input, mlir::MLIRContext *context) {
++    [](llvm::StringRef input, mlir::MLIRContext* context) {
+       if (printStablehloVersion.getValue()) {
+         auto version = stablehlo::getPortableArtifactVersion(input);
+         if (failed(version)) {
+@@ -346,14 +346,14 @@
+       }
+       return stablehlo::deserializePortableArtifact(input, context);
+     },
+-    [](DialectRegistry &registry) {
++    [](DialectRegistry& registry) {
+       mlir::registerAllDialects(registry);
+       mlir::stablehlo::registerAllDialects(registry);
+     });
+ 
+ }  //  namespace mlir
+ 
+-int main(int argc, char **argv) {
++int main(int argc, char** argv) {
+   return failed(
+       mlir::mlirTranslateMain(argc, argv, "StableHLO interpreter driver\n"));
+ }
+diff --ruN a/stablehlo/stablehlo/transforms/ChloDecompositionUtils.h b/stablehlo/stablehlo/transforms/ChloDecompositionUtils.h
+--- stablehlo/stablehlo/transforms/ChloDecompositionUtils.h
++++ stablehlo/stablehlo/transforms/ChloDecompositionUtils.h
+@@ -22,13 +22,13 @@
+ 
+ // Utility functions used in the Chlo to stablehlo legalization.
+ 
+-Value materializeLgamma(OpBuilder &rewriter, Location loc, ValueRange args);
++Value materializeLgamma(OpBuilder& rewriter, Location loc, ValueRange args);
+ 
+-Value materializeDigamma(OpBuilder &rewriter, Location loc, ValueRange args);
++Value materializeDigamma(OpBuilder& rewriter, Location loc, ValueRange args);
+ 
+-Value materializeZeta(OpBuilder &rewriter, Location loc, ValueRange args);
++Value materializeZeta(OpBuilder& rewriter, Location loc, ValueRange args);
+ 
+-Value materializePolygamma(OpBuilder &rewriter, Location loc, ValueRange args);
++Value materializePolygamma(OpBuilder& rewriter, Location loc, ValueRange args);
+ 
+ }  // namespace stablehlo
+ }  // namespace mlir
 diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp
 --- stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp
 +++ stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp
-@@ -2930,8 +2930,19 @@
+@@ -2930,8 +2930,18 @@
        SmallVector<int64_t> sliceShape(inputType.getShape().begin(),
                                        inputType.getShape().end());
        sliceShape[dim] = 1;
+-      auto sliceType =
+-          RankedTensorType::get(sliceShape, inputType.getElementType());
 +      Attribute sliceEncoding;
 +      ArrayRef<int64_t> inputBounds =
 +          hlo::encodingToBounds(inputType.getEncoding());
@@ -3578,14 +4963,12 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable
 +        sliceEncoding =
 +            hlo::boundsToEncoding(inputType.getEncoding(), sliceBounds);
 +      }
-       auto sliceType =
--          RankedTensorType::get(sliceShape, inputType.getElementType());
-+          RankedTensorType::get(sliceShape, inputType.getElementType(),
-+                                sliceEncoding);
++      auto sliceType = RankedTensorType::get(
++          sliceShape, inputType.getElementType(), sliceEncoding);
  
        Value slice;
        if (inputType.hasStaticShape()) {
-@@ -2980,9 +2991,18 @@
+@@ -2980,9 +2990,18 @@
        for (int64_t d = 0; d < rank; ++d) {
          if (d != dim) resultShape.push_back(sliceShape[d]);
        }
@@ -3605,7 +4988,7 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable
            slice);
        slicedInputs.push_back(reshaped);
      }
-@@ -3004,9 +3024,20 @@
+@@ -3004,9 +3023,20 @@
        SmallVector<int64_t> elementShape(outputType.getShape().begin(),
                                          outputType.getShape().end());
        elementShape[dim] = 1;
@@ -3627,4 +5010,207 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable
            newElement);
  
        int64_t outRank = outputType.getRank();
-
+diff --ruN a/stablehlo/stablehlo/transforms/PassPipelines.cpp b/stablehlo/stablehlo/transforms/PassPipelines.cpp
+--- stablehlo/stablehlo/transforms/PassPipelines.cpp
++++ stablehlo/stablehlo/transforms/PassPipelines.cpp
+@@ -24,7 +24,7 @@
+ namespace mlir {
+ namespace stablehlo {
+ 
+-void createStablehloDeserializePipeline(OpPassManager &pm) {
++void createStablehloDeserializePipeline(OpPassManager& pm) {
+   // Convert VHLO(version x.y.z) --> VHLO(current).
+   pm.addPass(stablehlo::createVhloToVersionPass(
+       {vhlo::Version::getCurrentVersion().toString()}));
+@@ -33,14 +33,14 @@
+   pm.addPass(stablehlo::createVhloLegalizeToStablehloPass());
+ }
+ 
+-void createChloPreSerializationPipeline(OpPassManager &pm) {
++void createChloPreSerializationPipeline(OpPassManager& pm) {
+   // CHLO --> StableHLO+Shape --> StableHLO
+   pm.addPass(stablehlo::createChloLegalizeToStablehloPass());
+   pm.addPass(stablehlo::createShapeLegalizeToStablehloPass());
+   pm.addPass(mlir::createReconcileUnrealizedCastsPass());
+ }
+ 
+-void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
++void createStablehloRemoveDynamismPipeline(OpPassManager& pm,
+                                            TypeRange refinedTypes) {
+   pm.addPass(stablehlo::createStablehloRefineArgumentsPass(refinedTypes));
+   pm.addPass(stablehlo::createStablehloRefineShapesPass());
+@@ -48,7 +48,7 @@
+       stablehlo::createStablehloCanonicalizeDynamismPass());
+ }
+ 
+-void createStablehloLowerQuantPipeline(OpPassManager &pm) {
++void createStablehloLowerQuantPipeline(OpPassManager& pm) {
+   pm.addNestedPass<mlir::func::FuncOp>(
+       stablehlo::createStablehloLegalizeQuantToMathPass());
+   pm.addNestedPass<mlir::func::FuncOp>(
+diff --ruN a/stablehlo/stablehlo/transforms/PassUtils.h b/stablehlo/stablehlo/transforms/PassUtils.h
+--- stablehlo/stablehlo/transforms/PassUtils.h
++++ stablehlo/stablehlo/transforms/PassUtils.h
+@@ -33,7 +33,7 @@
+ // Utility functions common across passes.
+ 
+ template <typename T>
+-Attribute getScalarLike(OpBuilder &b, T constant, Type type) {
++Attribute getScalarLike(OpBuilder& b, T constant, Type type) {
+   Type element = getElementTypeOrSelf(type);
+   if (isa<IntegerType>(element)) return b.getIntegerAttr(element, constant);
+   if (isa<FloatType>(element)) return b.getFloatAttr(element, constant);
+@@ -47,13 +47,13 @@
+ // in `scalar`.
+ // Returns stablehlo::ConstantOp if value type if static, else returns
+ // chlo::ConstantLikeOp.
+-Value getConstantLikeImpl(OpBuilder &b, Location loc, Attribute scalar,
++Value getConstantLikeImpl(OpBuilder& b, Location loc, Attribute scalar,
+                           Value val);
+ 
+ // Creates a chlo::ConstantLikeOp using a splat `constant` of the same shape
+ // as `val`.
+ template <typename T>
+-Value getConstantLike(OpBuilder &b, Location loc, T constant, Value val) {
++Value getConstantLike(OpBuilder& b, Location loc, T constant, Value val) {
+   auto shapedTy = cast<ShapedType>(val.getType());
+   Attribute scalar = getScalarLike(b, constant, shapedTy);
+   return getConstantLikeImpl(b, loc, scalar, val);
+@@ -63,7 +63,7 @@
+ // same shape as `val`.
+ // The distinction between double and APFloat causes issues so need this
+ // explicit template specialization.
+-Value getConstantLike(OpBuilder &b, Location loc, const APFloat &constant,
++Value getConstantLike(OpBuilder& b, Location loc, const APFloat& constant,
+                       Value val);
+ 
+ // Check if any of the given types are mlir::quant::QuantizedType.
+diff --ruN a/stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp b/stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp
+--- stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp
++++ stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp
+@@ -24,14 +24,14 @@
+ 
+ namespace {
+ 
+-static Value getConstantLikeMaxFiniteValue(OpBuilder &b, Location loc,
++static Value getConstantLikeMaxFiniteValue(OpBuilder& b, Location loc,
+                                            Value val) {
+   auto ty = cast<FloatType>(getElementTypeOrSelf(val.getType()));
+   return getConstantLike(
+       b, loc, llvm::APFloat::getLargest(ty.getFloatSemantics()), val);
+ }
+ 
+-static Value getConstantLikeInfValue(OpBuilder &b, Location loc, Value val,
++static Value getConstantLikeInfValue(OpBuilder& b, Location loc, Value val,
+                                      bool negative) {
+   auto ty = cast<FloatType>(getElementTypeOrSelf(val.getType()));
+   return getConstantLike(
+@@ -50,7 +50,7 @@
+             StablehloComplexMathExpanderPass>() {}
+ 
+  public:
+-  LogicalResult initialize(MLIRContext *context) override {
++  LogicalResult initialize(MLIRContext* context) override {
+     config.setUseTopDownTraversal(true);
+     RewritePatternSet patterns_(context);
+     populateStablehloComplexMathExpanderPatterns(context, &patterns_);
+@@ -76,8 +76,8 @@
+ 
+ }  // namespace
+ 
+-void populateStablehloComplexMathExpanderPatterns(MLIRContext *context,
+-                                                  RewritePatternSet *patterns) {
++void populateStablehloComplexMathExpanderPatterns(MLIRContext* context,
++                                                  RewritePatternSet* patterns) {
+   populateWithGenerated(*patterns);
+ }
+ 
+diff --ruN a/stablehlo/stablehlo/transforms/optimization/Passes.h b/stablehlo/stablehlo/transforms/optimization/Passes.h
+--- stablehlo/stablehlo/transforms/optimization/Passes.h
++++ stablehlo/stablehlo/transforms/optimization/Passes.h
+@@ -47,33 +47,33 @@
+ 
+ /// Collection of canonicalization patterns for StableHLO.
+ void populateStablehloCanonicalizationPatterns(
+-    MLIRContext *context, RewritePatternSet *patterns,
+-    const StablehloAggressiveSimplificationPassOptions &options,
++    MLIRContext* context, RewritePatternSet* patterns,
++    const StablehloAggressiveSimplificationPassOptions& options,
+     PatternBenefit benefit = 1);
+ 
+-void populateStablehloCanonicalizationPatterns(MLIRContext *context,
+-                                               RewritePatternSet *patterns,
++void populateStablehloCanonicalizationPatterns(MLIRContext* context,
++                                               RewritePatternSet* patterns,
+                                                PatternBenefit benefit = 1);
+ 
+ /// Collection of folding patterns for StableHLO.
+ void populateStablehloAggressiveFolderPatterns(
+-    MLIRContext *context, RewritePatternSet *patterns,
+-    const StablehloAggressiveFolderPassOptions &options,
++    MLIRContext* context, RewritePatternSet* patterns,
++    const StablehloAggressiveFolderPassOptions& options,
+     PatternBenefit benefit = 1);
+ 
+-void populateStablehloAggressiveFolderPatterns(MLIRContext *context,
+-                                               RewritePatternSet *patterns,
++void populateStablehloAggressiveFolderPatterns(MLIRContext* context,
++                                               RewritePatternSet* patterns,
+                                                PatternBenefit benefit = 1);
+ 
+ /// A subset of folding patterns for StableHLO that is necessary for shape
+ /// refinement.
+ void populateStablehloShapeFolderPatterns(
+-    MLIRContext *context, RewritePatternSet *patterns,
+-    const StablehloAggressiveFolderPassOptions &options,
++    MLIRContext* context, RewritePatternSet* patterns,
++    const StablehloAggressiveFolderPassOptions& options,
+     PatternBenefit benefit = 1);
+ 
+-void populateStablehloShapeFolderPatterns(MLIRContext *context,
+-                                          RewritePatternSet *patterns,
++void populateStablehloShapeFolderPatterns(MLIRContext* context,
++                                          RewritePatternSet* patterns,
+                                           PatternBenefit benefit = 1);
+ 
+ /// Some workloads in XLA import StableHLO from HLO. Since there are a few
+@@ -81,28 +81,28 @@
+ /// set of patterns brings the imported HLO back to a more canonical form
+ /// without applying a full set of graph simplifications.
+ void populateStablehloHloImportCanonicalizationPatterns(
+-    MLIRContext *context, RewritePatternSet *patterns,
+-    const StablehloAggressiveSimplificationPassOptions &options);
++    MLIRContext* context, RewritePatternSet* patterns,
++    const StablehloAggressiveSimplificationPassOptions& options);
+ 
+ void populateStablehloHloImportCanonicalizationPatterns(
+-    MLIRContext *context, RewritePatternSet *patterns);
++    MLIRContext* context, RewritePatternSet* patterns);
+ 
+ // If the `options` argument is provided, ensure it's not a temporary.
+ void populateStablehloCanonicalizationPatterns(
+-    MLIRContext *context, RewritePatternSet *patterns,
+-    StablehloAggressiveSimplificationPassOptions &&options,
++    MLIRContext* context, RewritePatternSet* patterns,
++    StablehloAggressiveSimplificationPassOptions&& options,
+     PatternBenefit benefit = 1) = delete;
+ void populateStablehloShapeFolderPatterns(
+-    MLIRContext *context, RewritePatternSet *patterns,
+-    StablehloAggressiveFolderPassOptions &&options,
++    MLIRContext* context, RewritePatternSet* patterns,
++    StablehloAggressiveFolderPassOptions&& options,
+     PatternBenefit benefit = 1) = delete;
+ void populateStablehloAggressiveFolderPatterns(
+-    MLIRContext *context, RewritePatternSet *patterns,
+-    StablehloAggressiveFolderPassOptions &&options,
++    MLIRContext* context, RewritePatternSet* patterns,
++    StablehloAggressiveFolderPassOptions&& options,
+     PatternBenefit benefit = 1) = delete;
+ void populateStablehloHloImportCanonicalizationPatterns(
+-    MLIRContext *context, RewritePatternSet *patterns,
+-    StablehloAggressiveSimplificationPassOptions &&options) = delete;
++    MLIRContext* context, RewritePatternSet* patterns,
++    StablehloAggressiveSimplificationPassOptions&& options) = delete;
+ 
+ }  // namespace stablehlo
+ }  // namespace mlir

From 07e23262d75736ce2340b7cf1d4e2eac6bcad0fa Mon Sep 17 00:00:00 2001
From: Majid Dadashi <majiddadashi@google.com>
Date: Fri, 18 Sep 2026 14:42:06 -0700
Subject: [PATCH 08/13] Fix the ping-pong bug in OptimizeBatchMatmulPass

The FC -> BMM was looking at T+FC to change that to BMM hoping that if a const + T is already folded by the time it sees it. However, resource + T is not immediately folded by design. So, it was leading to FC -> BMM -> FC -> BMM -> ...

This explicitly ensures a FC -> BMM doesn't kick in if rhs is from an eventually foldable chain.

PiperOrigin-RevId: 984076078
---
 .../lite/tests/optimize_batch_matmul.mlir     | 33 +++++++++++++++++++
 .../transforms/optimize_batch_matmul_pass.cc  |  9 +++--
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir b/tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir
index a65cad55124884..d06ce02b70d299 100644
--- a/tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir
@@ -242,4 +242,37 @@ func.func @FuseBatchMatmulToTransposeNegative(%arg0: tensor<2048x32x1x8x2xf32>,
   %196 = "tfl.batch_matmul"(%arg1, %195) <{adj_x = false, adj_y = false, asymmetric_quantize_inputs = false}> : (tensor<2048x1x32x2xf32>, tensor<2048x1x2x256xf32>) -> tensor<2048x1x32x256xf32>
   return %196 : tensor<2048x1x32x256xf32>
   // CHECK: "tfl.transpose"
+}
+
+// CHECK-LABEL: NotFuseTransposeFCRhsToBatchMatmulResourceConstant
+func.func @NotFuseTransposeFCRhsToBatchMatmulResourceConstant(%arg0: tensor<16x1024xf32>, %arg1: none) -> tensor<16x128xf32> {
+  %cst = arith.constant dense<[1, 0]> : tensor<2xi32>
+  %weight = arith.constant dense_resource<__elided__> : tensor<1024x128xf32>
+  %0 = "tfl.transpose"(%weight, %cst) : (tensor<1024x128xf32>, tensor<2xi32>) -> tensor<128x1024xf32>
+  // CHECK: "tfl.fully_connected"
+  // CHECK-NOT: "tfl.batch_matmul"
+  %1 = "tfl.fully_connected"(%arg0, %0, %arg1) {asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<16x1024xf32>, tensor<128x1024xf32>, none) -> tensor<16x128xf32>
+  func.return %1 : tensor<16x128xf32>
+}
+
+// CHECK-LABEL: NotFuseTransposeFCRhsToBatchMatmulReshapedResourceConstant
+func.func @NotFuseTransposeFCRhsToBatchMatmulReshapedResourceConstant(%arg0: tensor<16x1024xf32>, %arg1: none) -> tensor<16x128xf32> {
+  %cst = arith.constant dense<[1, 0]> : tensor<2xi32>
+  %cst_shape = arith.constant dense<[1024, 128]> : tensor<2xi32>
+  %weight = arith.constant dense_resource<__elided__> : tensor<1024x128xf32>
+  %0 = "tfl.reshape"(%weight, %cst_shape) : (tensor<1024x128xf32>, tensor<2xi32>) -> tensor<1024x128xf32>
+  %1 = "tfl.transpose"(%0, %cst) : (tensor<1024x128xf32>, tensor<2xi32>) -> tensor<128x1024xf32>
+  // CHECK: "tfl.fully_connected"
+  // CHECK-NOT: "tfl.batch_matmul"
+  %2 = "tfl.fully_connected"(%arg0, %1, %arg1) {asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<16x1024xf32>, tensor<128x1024xf32>, none) -> tensor<16x128xf32>
+  func.return %2 : tensor<16x128xf32>
+}
+
+// CHECK-LABEL: BatchMatmulResourceConstantToFullyConnected
+func.func @BatchMatmulResourceConstantToFullyConnected(%arg0: tensor<16x1024xf32>) -> tensor<16x128xf32> {
+  %weight = arith.constant dense_resource<__elided__> : tensor<1024x128xf32>
+  // CHECK: "tfl.fully_connected"
+  // CHECK-NOT: "tfl.batch_matmul"
+  %0 = "tfl.batch_matmul"(%arg0, %weight) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32>
+  func.return %0 : tensor<16x128xf32>
 }
\ No newline at end of file
diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc
index 31943b04681025..9bcda8d49a1902 100644
--- a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc
+++ b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
 #include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
 #include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
 #include "mlir/IR/PatternMatch.h"  // from @llvm-project
 #include "mlir/IR/TypeUtilities.h"  // from @llvm-project
 #include "mlir/IR/Value.h"  // from @llvm-project
@@ -50,7 +51,9 @@ bool NotFromFoldableChain(mlir::Value value) {
   mlir::Operation* defining_op = value.getDefiningOp();
 
   while (defining_op) {
-    if (mlir::isa<DequantizeOp>(defining_op)) {
+    if (mlir::isa<DequantizeOp>(defining_op) ||
+        defining_op->hasTrait<mlir::OpTrait::ConstantLike>() ||
+        matchPattern(defining_op, mlir::m_Constant())) {
       return false;
     }
 
@@ -59,8 +62,10 @@ bool NotFromFoldableChain(mlir::Value value) {
       defining_op = reshape_op.getInput().getDefiningOp();
     } else if (auto split_op = mlir::dyn_cast<SplitOp>(defining_op)) {
       defining_op = split_op.getValue().getDefiningOp();
+    } else if (auto transpose_op = mlir::dyn_cast<TransposeOp>(defining_op)) {
+      defining_op = transpose_op.getInput().getDefiningOp();
     } else {
-      // Stop if the op is not Dequantize, Reshape, or Split.
+      // Stop if the op does not preserve constant nature.
       break;
     }
   }

From 2e6f1ef8cfaaa13fe2eb0bb0c90027facc0cac57 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 18 Sep 2026 15:12:54 -0700
Subject: [PATCH 09/13] Increase shard_count of np_indexing_test to 30 to fix
 TSAN timeouts.

PiperOrigin-RevId: 984092201
---
 tensorflow/python/ops/numpy_ops/tests/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/ops/numpy_ops/tests/BUILD b/tensorflow/python/ops/numpy_ops/tests/BUILD
index f7d7e47bd8dc91..d6bb0889655e3d 100644
--- a/tensorflow/python/ops/numpy_ops/tests/BUILD
+++ b/tensorflow/python/ops/numpy_ops/tests/BUILD
@@ -252,7 +252,7 @@ py_test(
         "--num_generated_cases=90",
         "--enable_x64",  # Needed to enable dtype check
     ],
-    shard_count = 10,
+    shard_count = 30,
     strict_deps = True,
     # TODO(b/164245103): Re-enable GPU once tf.tensor_strided_slice_update's segfault is fixed.
     tags = [

From 720136c7d69b1df4d54c43d5c9937a5132ccb3b4 Mon Sep 17 00:00:00 2001
From: Junwhan Ahn <junwhan@google.com>
Date: Fri, 18 Sep 2026 15:34:08 -0700
Subject: [PATCH 10/13] Round the device memory size to the nearest GiB while
 producing topology fingerprint

PiperOrigin-RevId: 984102296
---
 .../pjrt/gpu/se_gpu_topology_description.cc   | 21 ++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc
index 020d669b9f9e3d..a4cb5769b865ea 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc
@@ -170,8 +170,27 @@ StreamExecutorGpuTopologyDescription::CreateDeviceDescription(
 
 absl::StatusOr<uint64_t> StreamExecutorGpuTopologyDescription::Fingerprint()
     const {
+  GpuTopologyProto proto = gpu_topology_->ToProto();
+
+  stream_executor::GpuDeviceInfoProto* device_info =
+      proto.mutable_gpu_target_config()->mutable_gpu_device_info();
+  if (device_info->device_memory_size() > 0) {
+    // Round `device_memory_size` to the nearest GiB to produce deterministic
+    // fingerprint even if the memory size is slightly different.
+    //
+    // TODO(b/563487743): Fix `device_memory_size` to return a stable lower
+    // bound instead and get rid of this logic.
+    static constexpr int64_t kGiB = 1 << 30;
+    device_info->set_device_memory_size(
+        (device_info->device_memory_size() + kGiB / 2) / kGiB * kGiB);
+  }
+  // Exclude `model_str` from the fingerprint as the rest of the proto already
+  // contains the same info and the device memory size in model str can make the
+  // fingerprint non-portable across GPUs.
+  device_info->clear_model_str();
+
   std::string result;
-  if (!tsl::SerializeToStringDeterministic(gpu_topology_->ToProto(), &result)) {
+  if (!tsl::SerializeToStringDeterministic(proto, &result)) {
     return absl::InternalError("Failed to serialize gpu_topology");
   }
   return tsl::Fingerprint64(result);

From 10b878183f3958bf67510eb6058cef16eccd6500 Mon Sep 17 00:00:00 2001
From: Enver Kayaaslan <enver@google.com>
Date: Fri, 18 Sep 2026 15:40:26 -0700
Subject: [PATCH 11/13] Move USE_MATH_DEFINES macro above the includes on hlo
 evaluator.

PiperOrigin-RevId: 984105430
---
 third_party/xla/xla/hlo/evaluator/hlo_evaluator.h | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h
index b6fd3bad7c2515..9f2222ed9d029f 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h
@@ -16,14 +16,12 @@ limitations under the License.
 #ifndef XLA_HLO_EVALUATOR_HLO_EVALUATOR_H_
 #define XLA_HLO_EVALUATOR_HLO_EVALUATOR_H_
 
-#include "absl/log/log.h"
-#include "absl/status/status_macros.h"
 #define _USE_MATH_DEFINES
 
 #include <complex>
-#include <cstddef>  // NOLINT(build/include_order)
+#include <cstddef>
 #include <cstdint>
-#include <deque>  // NOLINT(build/include_order)
+#include <deque>
 #include <functional>
 #include <memory>
 #include <optional>
@@ -35,7 +33,9 @@ limitations under the License.
 #include "absl/container/inlined_vector.h"
 #include "absl/container/node_hash_map.h"
 #include "absl/log/check.h"
+#include "absl/log/log.h"
 #include "absl/status/status.h"
+#include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"

From 2ab8f03dacadedf056b81db0a958280118b74985 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 18 Sep 2026 17:00:16 -0700
Subject: [PATCH 12/13] Before CheckLayouts(), update computation_layouts_ with
 latest entry computation layout.

PiperOrigin-RevId: 984141229
---
 third_party/xla/xla/service/layout_assignment.cc | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/third_party/xla/xla/service/layout_assignment.cc b/third_party/xla/xla/service/layout_assignment.cc
index 6ba874a592a29e..425e2f46e154bd 100644
--- a/third_party/xla/xla/service/layout_assignment.cc
+++ b/third_party/xla/xla/service/layout_assignment.cc
@@ -3738,6 +3738,12 @@ absl::StatusOr<bool> LayoutAssignment::RunImpl(
                                               entry_computation_layout_));
 
 #ifndef NDEBUG
+  // Before validation, update computation_layouts_ with the entry computation
+  // layout since it might be modified during PropagateComputationLayouts()
+  // above for AUTO layouts.
+  *mutable_computation_constraints(module->entry_computation())
+       ->mutable_computation_constraint()
+       ->mutable_computation_layout() = *entry_computation_layout_;
   ABSL_RETURN_IF_ERROR(CheckLayouts(module, execution_threads));
 #endif  // NDEBUG
 

From 62ffa09a8f6ebeec83671e21a0c4267ab7641a9d Mon Sep 17 00:00:00 2001
From: "Patrick C. Toulme" <ptoulme@google.com>
Date: Fri, 18 Sep 2026 17:05:28 -0700
Subject: [PATCH 13/13] Copy removal: count a view's readers as uses of the
 viewed value

CopyInsertion takes an optional view color. A view colored value is an
address stand in that owns no storage and aliases its operand 0's buffer.
Copy removal then counts the transitive readers of a value's view colored
users as uses of that value, so no producer is merged into the viewed
buffer between a view and its readers. Backends without views are
unchanged.

PiperOrigin-RevId: 984143462
---
 third_party/xla/xla/service/BUILD             |   6 +
 third_party/xla/xla/service/copy_insertion.cc |   2 +-
 third_party/xla/xla/service/copy_insertion.h  |  15 +-
 .../xla/xla/service/copy_insertion_test.cc    | 178 ++++++++++++++++++
 third_party/xla/xla/service/copy_removal.cc   |  90 +++++++--
 third_party/xla/xla/service/copy_removal.h    |  43 ++++-
 6 files changed, 318 insertions(+), 16 deletions(-)

diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index ce553f11f86a9d..9fd8b08de32c83 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -3657,6 +3657,8 @@ xla_cc_test(
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
         "//xla/hlo/analysis:alias_info",
+        "//xla/hlo/analysis:hlo_alias_analysis",
+        "//xla/hlo/analysis:hlo_ordering",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/parser:hlo_parser",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
@@ -3670,6 +3672,10 @@ xla_cc_test(
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_benchmark//:benchmark",
         "@com_google_googletest//:gtest",
diff --git a/third_party/xla/xla/service/copy_insertion.cc b/third_party/xla/xla/service/copy_insertion.cc
index 1f6dfbb032fb16..b9ca9d683cdfa6 100644
--- a/third_party/xla/xla/service/copy_insertion.cc
+++ b/third_party/xla/xla/service/copy_insertion.cc
@@ -1821,7 +1821,7 @@ absl::Status CopyInsertion::RemoveUnnecessaryCopies(
   ABSL_ASSIGN_OR_RETURN(std::unique_ptr<HloAliasAnalysis> alias_analysis,
                    HloAliasAnalysis::Run(module, alias_info_));
   CopyRemover copy_remover(*module, *alias_analysis, alias_info_,
-                           ordering.get(), execution_threads);
+                           ordering.get(), execution_threads, view_color_);
   if (VLOG_IS_ON(3)) {
     LOG(INFO) << "Removing unnecessary copies in " << module->name();
     LOG(INFO) << "Buffer values, in dependency order: ";
diff --git a/third_party/xla/xla/service/copy_insertion.h b/third_party/xla/xla/service/copy_insertion.h
index 9b84ced52c75e3..135a7b94cde138 100644
--- a/third_party/xla/xla/service/copy_insertion.h
+++ b/third_party/xla/xla/service/copy_insertion.h
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include <cstdint>
 #include <functional>
+#include <optional>
 
 #include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"
@@ -62,17 +63,26 @@ class CopyInsertion : public HloModulePass {
   // backend specific function that decides whether an instruction
   // can share buffer with its operand.
   //
+  // `view_color` is the layout memory space color of view values (address
+  // stand ins with no storage of their own that alias their operand 0's
+  // buffer, see BufferAssigner::Options::dus_view_color). When set, copy
+  // removal counts the readers of a value's views as uses of that value, so
+  // no producer is merged into the viewed buffer between a view and its
+  // readers.
+  //
   // TODO(b/80315712): Find a better way to tell whether a fusion can share
   // buffer.
   explicit CopyInsertion(
       const AliasInfo* alias_info,
       int64_t use_region_based_live_range_analysis = kUseRegionAnalysisLimit,
       std::function<bool(const HloInstruction* copy)> should_skip_removal =
-          nullptr)
+          nullptr,
+      std::optional<int64_t> view_color = std::nullopt)
       : alias_info_(alias_info),
         should_skip_removal_(should_skip_removal),
         use_region_based_live_range_analysis_(
-            use_region_based_live_range_analysis) {}
+            use_region_based_live_range_analysis),
+        view_color_(view_color) {}
 
   // Try to remove as many copies from the module as possible without
   // introducing live range interference. Only copy instructions that are
@@ -140,6 +150,7 @@ class CopyInsertion : public HloModulePass {
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads);
   int64_t use_region_based_live_range_analysis_;
+  std::optional<int64_t> view_color_;
 };
 
 }  // namespace xla
diff --git a/third_party/xla/xla/service/copy_insertion_test.cc b/third_party/xla/xla/service/copy_insertion_test.cc
index fdac453863b75d..ce2da447ae0a58 100644
--- a/third_party/xla/xla/service/copy_insertion_test.cc
+++ b/third_party/xla/xla/service/copy_insertion_test.cc
@@ -15,8 +15,10 @@ limitations under the License.
 
 #include "xla/service/copy_insertion.h"
 
+#include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <string>
 #include <tuple>
 #include <vector>
@@ -26,11 +28,20 @@ limitations under the License.
 #include "absl/algorithm/container.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
+#include "absl/status/status_macros.h"
+#include "absl/status/status_matchers.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_split.h"
 #include "absl/strings/string_view.h"
+#include "absl/strings/substitute.h"
 #include "xla/comparison_util.h"
 #include "xla/debug_options_flags.h"
 #include "xla/frontend_attributes.h"
 #include "xla/hlo/analysis/alias_info.h"
+#include "xla/hlo/analysis/hlo_alias_analysis.h"
+#include "xla/hlo/analysis/hlo_ordering.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
@@ -46,6 +57,7 @@ limitations under the License.
 #include "xla/layout_util.h"
 #include "xla/literal_util.h"
 #include "xla/service/buffer_value.h"
+#include "xla/service/copy_removal.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
@@ -3918,6 +3930,172 @@ ENTRY entry {
   EXPECT_THAT(while_1, op::While(op::Tuple(op::Copy())));
 }
 
+// A view (a custom call colored `kViewColor` that addresses into its operand
+// 0's buffer without storage of its own) is read by its readers, not at the
+// view. With `view_color` set, copy removal counts those readers as uses of
+// the viewed value: the loop carried stack's root copy protects the elementwise
+// producer that would otherwise be merged into the stack's buffer between the
+// view and its reader ($0 orders the reader after the producer). Without the
+// option, or with the reader ordered before the producer, the copy is elided.
+TEST_F(CopyInsertionTest, ViewReadersCountAsUsesOfTheViewedValue) {
+  constexpr int64_t kViewColor = 5;
+  constexpr absl::string_view kHloTemplate = R"(
+HloModule view_readers, is_scheduled=true
+
+cond {
+  p = (s32[], f32[4,8], f32[1,8]) parameter(0)
+  i = s32[] get-tuple-element(p), index=0
+  n = s32[] constant(3)
+  ROOT lt = pred[] compare(i, n), direction=LT
+}
+
+body {
+  p = (s32[], f32[4,8], f32[1,8]) parameter(0)
+  i = s32[] get-tuple-element(p), index=0
+  stack = f32[4,8] get-tuple-element(p), index=1
+  acc = f32[1,8] get-tuple-element(p), index=2
+  c0 = s32[] constant(0)
+  view = f32[1,8]{1,0:S(5)} custom-call(stack, i, c0), custom_call_target="view"
+  one = f32[] constant(1)
+  ones = f32[4,8] broadcast(one), dimensions={}
+  $0
+  acc_new = f32[1,8] add(acc, reader)
+  one_s32 = s32[] constant(1)
+  i_new = s32[] add(i, one_s32)
+  copy = f32[4,8] copy(new_stack)
+  ROOT t = (s32[], f32[4,8], f32[1,8]) tuple(i_new, copy, acc_new)
+}
+
+ENTRY entry {
+  stack0 = f32[4,8] parameter(0)
+  acc0 = f32[1,8] parameter(1)
+  i0 = s32[] constant(0)
+  init = (s32[], f32[4,8], f32[1,8]) tuple(i0, stack0, acc0)
+  ROOT loop = (s32[], f32[4,8], f32[1,8]) while(init), condition=cond, body=body
+}
+)";
+  constexpr absl::string_view kReaderAfterProducer =
+      "new_stack = f32[4,8] add(stack, ones)\n"
+      "  reader = f32[1,8] negate(view)";
+  constexpr absl::string_view kReaderBeforeProducer =
+      "reader = f32[1,8] negate(view)\n"
+      "  new_stack = f32[4,8] add(stack, ones)";
+  auto root_copy_survives =
+      [&](absl::string_view order,
+          std::optional<int64_t> view_color) -> absl::StatusOr<bool> {
+    ABSL_ASSIGN_OR_RETURN(
+        std::unique_ptr<HloModule> module,
+        ParseAndReturnVerifiedModule(absl::Substitute(kHloTemplate, order)));
+    CopyInsertion copy_insertion(&alias_info_,
+                                 /*use_region_based_live_range_analysis=*/-1,
+                                 /*should_skip_removal=*/nullptr, view_color);
+    ABSL_RETURN_IF_ERROR(copy_insertion.RemoveUnnecessaryCopies(module.get()));
+    const HloInstruction* root =
+        module->GetComputationWithName("body")->root_instruction();
+    return root->operand(1)->opcode() == HloOpcode::kCopy;
+  };
+  EXPECT_THAT(root_copy_survives(kReaderAfterProducer, kViewColor),
+              absl_testing::IsOkAndHolds(true));
+  EXPECT_THAT(root_copy_survives(kReaderAfterProducer, std::nullopt),
+              absl_testing::IsOkAndHolds(false));
+  EXPECT_THAT(root_copy_survives(kReaderBeforeProducer, kViewColor),
+              absl_testing::IsOkAndHolds(false));
+}
+
+// A view colored user that writes through the view (an in place op aliasing
+// its output onto the view operand) is a reader at its own position, not a
+// forwarder: the copies copy insertion places around it are elided exactly
+// as without a view color, and the in place writer lands in the viewed
+// buffer.
+TEST_F(CopyInsertionTest, ViewWriterIsNotForwardedAsAView) {
+  constexpr int64_t kViewColor = 5;
+  constexpr absl::string_view kHlo = R"(
+HloModule view_writer, is_scheduled=true
+
+ENTRY entry {
+  base = f32[4,8] parameter(0)
+  update = f32[1,8] parameter(1)
+  c0 = s32[] constant(0)
+  c1 = s32[] constant(1)
+  view = f32[1,8]{1,0:S(5)} custom-call(base, c1, c0), custom_call_target="view"
+  view_copy = f32[1,8]{1,0:S(5)} copy(view)
+  writer = f32[1,8]{1,0:S(5)} custom-call(view_copy, update), custom_call_target="write_through_view", output_to_operand_aliasing={{}: (0, {})}
+  base_copy = f32[4,8] copy(base)
+  ROOT dus = f32[4,8] dynamic-update-slice(base_copy, writer, c1, c0)
+}
+)";
+  for (std::optional<int64_t> view_color :
+       {std::optional<int64_t>(kViewColor), std::optional<int64_t>()}) {
+    ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                         ParseAndReturnVerifiedModule(kHlo));
+    CopyInsertion copy_insertion(&alias_info_,
+                                 /*use_region_based_live_range_analysis=*/-1,
+                                 /*should_skip_removal=*/nullptr, view_color);
+    ASSERT_OK(copy_insertion.RemoveUnnecessaryCopies(module.get()));
+    const HloInstruction* dus = module->entry_computation()->root_instruction();
+    EXPECT_EQ(dus->operand(0)->name(), "base")
+        << "view color " << view_color.value_or(-1) << ":\n"
+        << module->ToString();
+    EXPECT_EQ(dus->operand(1)->operand(0)->name(), "view")
+        << "view color " << view_color.value_or(-1) << ":\n"
+        << module->ToString();
+  }
+}
+
+// The synthesized view uses are seeded from the viewed value only. A view's
+// own readers are already its dataflow uses (through its bitcast position),
+// so walking the view's view colored bitcast again would list every reader a
+// second time and inflate the region analysis use counts.
+TEST_F(CopyInsertionTest, ViewReadersAreListedOncePerValue) {
+  constexpr int64_t kViewColor = 5;
+  constexpr absl::string_view kHlo = R"(
+HloModule view_readers_once, is_scheduled=true
+
+ENTRY entry {
+  stack = f32[4,8] parameter(0)
+  i = s32[] parameter(1)
+  c0 = s32[] constant(0)
+  view = f32[1,8]{1,0:S(5)} custom-call(stack, i, c0), custom_call_target="view"
+  row = f32[8]{0:S(5)} bitcast(view)
+  reader = f32[8] negate(row)
+  other = f32[8] exponential(row)
+  ROOT t = (f32[8], f32[8]) tuple(reader, other)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHlo));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloAliasAnalysis> alias_analysis,
+                       HloAliasAnalysis::Run(module.get(), &alias_info_));
+  DependencyHloOrdering ordering(module.get());
+  CopyRemover remover(*module, *alias_analysis, &alias_info_, &ordering,
+                      /*execution_threads=*/{}, kViewColor);
+  // One "<id> <name>, uses: ..." line per value in CopyRemover::ToString.
+  auto uses_line_of = [&](absl::string_view name) -> std::string {
+    for (absl::string_view line : absl::StrSplit(remover.ToString(), '\n')) {
+      if (absl::StrContains(line, absl::StrCat(" ", name, ">, uses: "))) {
+        return std::string(line);
+      }
+    }
+    return "";
+  };
+  auto count = [](absl::string_view haystack, absl::string_view needle) {
+    int n = 0;
+    for (size_t pos = haystack.find(needle); pos != absl::string_view::npos;
+         pos = haystack.find(needle, pos + needle.size())) {
+      ++n;
+    }
+    return n;
+  };
+  const std::string stack_uses = uses_line_of("stack");
+  const std::string view_uses = uses_line_of("view");
+  ASSERT_FALSE(stack_uses.empty()) << remover.ToString();
+  ASSERT_FALSE(view_uses.empty()) << remover.ToString();
+  for (absl::string_view reader : {"reader, operand 0", "other, operand 0"}) {
+    EXPECT_EQ(count(stack_uses, reader), 1) << stack_uses;
+    EXPECT_EQ(count(view_uses, reader), 1) << view_uses;
+  }
+}
+
 TEST_F(CopyInsertionTest, InPlaceCollectivePermuteCopy) {
   absl::string_view hlo_string = R"(
 HloModule hlo_runner_test_0.1
diff --git a/third_party/xla/xla/service/copy_removal.cc b/third_party/xla/xla/service/copy_removal.cc
index e968cffed95399..123df19612165f 100644
--- a/third_party/xla/xla/service/copy_removal.cc
+++ b/third_party/xla/xla/service/copy_removal.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <cstdint>
 #include <functional>
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -624,10 +625,12 @@ Relation::RuntimeOrder ComputeRelativeLocation::ComputeRuntimeOrdering(
 CopyRemover::CopyRemover(
     const HloModule& module, const HloAliasAnalysis& alias_analysis,
     const AliasInfo* alias_info, HloOrdering* ordering,
-    const absl::flat_hash_set<absl::string_view>& execution_threads)
+    const absl::flat_hash_set<absl::string_view>& execution_threads,
+    std::optional<int64_t> view_color)
     : dataflow_(alias_analysis.dataflow_analysis()),
       alias_info_(alias_info),
-      ordering_(ordering) {
+      ordering_(ordering),
+      view_color_(view_color) {
   // Instruction indices based on post order traversal of computations and
   // instructions. Used as an enhancement for getting strict weak ordering
   // used for sorting below.
@@ -792,6 +795,9 @@ void CopyRemover::AddValueList(
     for (const HloUse& use : value->GetUses()) {
       new_node->uses.push_back(&use);
     }
+    if (view_color_.has_value()) {
+      AddViewUses(value, new_node);
+    }
 
     // Connect the new node into the linked list.
     if (tail == nullptr) {
@@ -809,6 +815,63 @@ void CopyRemover::AddValueList(
   value_lists_.insert(head);
 }
 
+void CopyRemover::AddViewUses(const HloValue* value, ValueNode* node) {
+  auto is_view = [this](const HloInstruction* instruction) {
+    return instruction->shape().IsArray() &&
+           instruction->shape().has_layout() &&
+           instruction->shape().layout().memory_space() == *view_color_;
+  };
+  // A view's own readers are its dataflow uses already; seeding from its view
+  // colored bitcasts would list every reader a second time.
+  if (is_view(value->defining_instruction())) {
+    return;
+  }
+  // A view colored bitcast, or a view colored copy (an address copy until
+  // copy insertion elides it), forwards the address to its own users. Any
+  // other view colored user is recorded as a reader at its own position, on
+  // purpose: nothing here orders it against the other readers of the viewed
+  // value, so the pass that creates views must never let a reader write
+  // through one. An in place writer through a view is sound only when its
+  // write is also a dataflow use of the viewed buffer (for example the
+  // dynamic-update-slice it feeds), which orders it.
+  auto forwards_view = [&](const HloInstruction* instruction) {
+    return is_view(instruction) &&
+           (instruction->opcode() == HloOpcode::kBitcast ||
+            instruction->opcode() == HloOpcode::kCopy);
+  };
+  // Only the viewed value itself (operand 0 of the view) is read through the
+  // view; a view's other operands (start indices) are consumed at the view's
+  // own position.
+  absl::flat_hash_set<const HloInstruction*> visited;
+  std::vector<const HloInstruction*> worklist;
+  for (const HloUse& use : value->GetUses()) {
+    if (use.operand_number == 0 && is_view(use.instruction) &&
+        visited.insert(use.instruction).second) {
+      worklist.push_back(use.instruction);
+    }
+  }
+  while (!worklist.empty()) {
+    const HloInstruction* view = worklist.back();
+    worklist.pop_back();
+    for (HloInstruction* user : view->users()) {
+      if (forwards_view(user)) {
+        if (visited.insert(user).second) {
+          worklist.push_back(user);
+        }
+        continue;
+      }
+      for (int64_t i = 0; i < user->operand_count(); ++i) {
+        if (user->operand(i) != view) {
+          continue;
+        }
+        const HloUse& use = view_uses_.emplace_back(user, i, ShapeIndex{});
+        view_uses_by_pointer_.insert(&use);
+        node->uses.push_back(&use);
+      }
+    }
+  }
+}
+
 // This method also fills in copy_map_ which indicates which nodes
 // in the value lists corresponding to the source and destination values of
 // kCopy instructions. value_to_node should map each HloValue to its
@@ -859,8 +922,10 @@ absl::Status CopyRemover::Verify() const {
         TF_RET_CHECK(copy_map_.at(def).dest == p);
       }
       for (const HloUse* use : p->uses) {
+        // A copy reading a view of p's value is a view use of p, but its copy
+        // source is the view's value.
         if (use->instruction->opcode() == HloOpcode::kCopy &&
-            ContainsKey(copy_map_, use->instruction)) {
+            ContainsKey(copy_map_, use->instruction) && !IsViewUse(use)) {
           TF_RET_CHECK(copy_map_.at(use->instruction).src == p);
         }
       }
@@ -1299,22 +1364,25 @@ void CopyRemover::RemoveCopyValue(ValueNode* copy_value_node,
   copy_value_node->prev->next = copy_value_node->next;
   copy_value_node->next->prev = copy_value_node->prev;
 
-  // Patch up uses. Remove use of copy from operand_node uses.
-  auto it =
-      absl::c_find_if(operand_node->uses, [copy_value_node](const HloUse* use) {
-        return use->instruction ==
-               copy_value_node->value->defining_instruction();
-      });
+  // Patch up uses. Remove use of copy from operand_node uses. A view use at
+  // the same copy (the copy reads a view of the operand) stays: the copy still
+  // reads the operand's buffer there.
+  auto it = absl::c_find_if(operand_node->uses, [this, copy_value_node](
+                                                    const HloUse* use) {
+    return use->instruction == copy_value_node->value->defining_instruction() &&
+           !IsViewUse(use);
+  });
   CHECK(it != operand_node->uses.end());
   operand_node->uses.erase(it);
 
   // If the elided copy has any uses which are themselves kCopy instructions
   // then patch up the copy info to reflect the that this kCopy instruction
-  // has a different operand (the operand of the elided copy).
+  // has a different operand (the operand of the elided copy). A view use at
+  // a copy keeps its own source, the view's value.
   for (const HloUse* copy_use : copy_value_node->uses) {
     operand_node->uses.push_back(copy_use);
     if (copy_use->instruction->opcode() == HloOpcode::kCopy &&
-        ContainsKey(copy_map_, copy_use->instruction)) {
+        ContainsKey(copy_map_, copy_use->instruction) && !IsViewUse(copy_use)) {
       copy_map_.at(copy_use->instruction).src = operand_node;
     }
   }
diff --git a/third_party/xla/xla/service/copy_removal.h b/third_party/xla/xla/service/copy_removal.h
index 86b578c3276763..c135ca711b618d 100644
--- a/third_party/xla/xla/service/copy_removal.h
+++ b/third_party/xla/xla/service/copy_removal.h
@@ -17,7 +17,9 @@ limitations under the License.
 #define XLA_SERVICE_COPY_REMOVAL_H_
 
 #include <cstdint>
+#include <deque>
 #include <functional>
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -392,7 +394,8 @@ class CopyRemover {
 
     // The uses are maintained outside of HloValue::uses() because
     // HloValue::uses() is not updatable (a fully updatable dataflow analysis
-    // is slow).
+    // is slow). Besides the value's own uses this holds the uses of the
+    // readers of its views (see AddViewUses).
     std::vector<const HloUse*> uses;
 
     // next/prev elements in the linked list. The list is circularly linked so
@@ -401,9 +404,16 @@ class CopyRemover {
     ValueNode* next = nullptr;
   };
 
+  // `view_color` is the layout memory space color of view values: address
+  // stand ins that own no storage and alias the buffer of their operand 0
+  // (BufferAssigner::Options::dus_view_color). A reader of such a value reads
+  // the viewed buffer at its own position, so it counts as a use of the
+  // viewed value; otherwise a producer merged into that buffer between the
+  // view and its reader would be read through the view.
   CopyRemover(const HloModule& module, const HloAliasAnalysis& alias_analysis,
               const AliasInfo* alias_info, HloOrdering* ordering,
-              const absl::flat_hash_set<absl::string_view>& execution_threads);
+              const absl::flat_hash_set<absl::string_view>& execution_threads,
+              std::optional<int64_t> view_color = std::nullopt);
 
   // Add a list containing the given values to CopyRemover. This
   // represents the values contained in a single buffer. For each value in
@@ -413,6 +423,29 @@ class CopyRemover {
       absl::Span<const HloValue* const> values,
       absl::flat_hash_map<const HloValue*, ValueNode*>* value_to_node);
 
+  // Appends to `node->uses` a use at every reader of a view of `value`: a
+  // use at operand 0 of a view colored instruction, followed through the view
+  // colored bitcasts and copies that forward its address (the reader side of
+  // ExtendViewBaseLiveRanges in hlo_live_range.cc).
+  //
+  // The walk is flat, and correct only under a contract the pass that creates
+  // views enforces: a view colored value never escapes into a tuple, a called
+  // computation (while, call, conditional) or an async pair, where the reads
+  // would happen later than the position recorded here; a reader never
+  // writes through it; and its operand 0 is never itself view colored (views
+  // are one level deep, so the readers found here read `value`'s own buffer).
+  // A view colored user that is not a bitcast or copy is recorded as a reader
+  // at its own position on purpose: an in place writer through a view is
+  // ordered only when its write is also a dataflow use of the viewed buffer.
+  void AddViewUses(const HloValue* value, ValueNode* node);
+
+  // Returns true if `use` was synthesized by AddViewUses rather than taken
+  // from an HloValue: such a use reads the value through a view, not as the
+  // operand it names.
+  bool IsViewUse(const HloUse* use) const {
+    return view_uses_by_pointer_.contains(use);
+  }
+
   // This method also fills in copy_map_ which indicates which nodes
   // in the value lists corresponding to the source and destination values of
   // kCopy instructions. value_to_node should map each HloValue to its
@@ -537,6 +570,12 @@ class CopyRemover {
     ValueNode* dest = nullptr;
   };
   absl::flat_hash_map<const HloInstruction*, CopyNodes> copy_map_;
+
+  // See the constructor. The synthesized uses live here so ValueNode::uses can
+  // point at them like at HloValue uses; a deque keeps the pointers stable.
+  std::optional<int64_t> view_color_;
+  std::deque<HloUse> view_uses_;
+  absl::flat_hash_set<const HloUse*> view_uses_by_pointer_;
 };
 };  // namespace xla